Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 47 additions & 27 deletions src/mcp/mcp.c
Original file line number Diff line number Diff line change
Expand Up @@ -4195,18 +4195,19 @@ static char *bm25_search(cbm_store_t *store, const char *project, const char *qu
}

/* Extract keyword strings from a yyjson array into `keywords`. Returns the
* number of strings copied (capped at `max_out`). */
* number of strings copied (capped at `max_out`), or -1 when any element is
* not a string: a mixed-type array is a caller error and is never silently
* narrowed to its string members. */
static int extract_semantic_keywords(yyjson_val *sq_val, const char **keywords, int max_out) {
int kw_count = (int)yyjson_arr_size(sq_val);
if (kw_count > max_out) {
kw_count = max_out;
}
size_t kw_idx = 0;
size_t kw_max = 0;
yyjson_val *kw_val;
int ki = 0;
yyjson_arr_foreach(sq_val, kw_idx, kw_max, kw_val) {
if (ki < kw_count && yyjson_is_str(kw_val)) {
if (!yyjson_is_str(kw_val)) {
return -1;
}
if (ki < max_out) {
keywords[ki++] = yyjson_get_str(kw_val);
}
}
Expand Down Expand Up @@ -4237,13 +4238,20 @@ static void emit_semantic_results(yyjson_mut_doc *doc, yyjson_mut_val *root,
yyjson_mut_obj_add_val(doc, root, "semantic", sem);
}

typedef enum {
SQ_RUN_OK = 0,
SQ_RUN_TYPE_ERROR, /* semantic_query is not an array of strings */
SQ_RUN_STORE_ERROR, /* the vector scan itself failed */
} sq_run_status_t;

/* Run the semantic_query vector search from raw args. Sets *out_vresults /
* *out_vcount (caller frees via cbm_store_free_vector_results when vcount>0).
* Returns true if semantic_query was provided as a non-array (type error —
* caller should surface to the user). */
static bool run_semantic_query_core(const char *args, cbm_store_t *store, const char *project,
int materialize_limit, cbm_vector_result_t **out_vresults,
int *out_vcount, bool *out_present) {
* A store without a vector table (lean index) is an empty result; a failed
* scan is SQ_RUN_STORE_ERROR and must never be rendered as zero matches. */
static sq_run_status_t run_semantic_query_core(const char *args, cbm_store_t *store,
const char *project, int materialize_limit,
cbm_vector_result_t **out_vresults, int *out_vcount,
bool *out_present) {
enum { MAX_KW_SEARCH = 32 };
*out_vresults = NULL;
*out_vcount = 0;
Expand All @@ -4256,25 +4264,31 @@ static bool run_semantic_query_core(const char *args, cbm_store_t *store, const
if (out_present && sq_val) {
*out_present = true;
}
bool type_error = false;
sq_run_status_t status = SQ_RUN_OK;
if (sq_val && !yyjson_is_arr(sq_val)) {
type_error = true;
status = SQ_RUN_TYPE_ERROR;
} else if (sq_val && yyjson_arr_size(sq_val) > 0) {
const char *keywords[MAX_KW_SEARCH];
int ki = extract_semantic_keywords(sq_val, keywords, MAX_KW_SEARCH);
cbm_vector_result_t *vresults = NULL;
int vcount = 0;
if (cbm_store_vector_search(store, project, keywords, ki, materialize_limit, &vresults,
&vcount) == CBM_STORE_OK &&
vcount > 0) {
*out_vresults = vresults;
*out_vcount = vcount;
if (ki < 0) {
status = SQ_RUN_TYPE_ERROR;
} else {
cbm_vector_result_t *vresults = NULL;
int vcount = 0;
int rc = cbm_store_vector_search(store, project, keywords, ki, materialize_limit,
&vresults, &vcount);
if (rc == CBM_STORE_ERR) {
status = SQ_RUN_STORE_ERROR;
} else if (rc == CBM_STORE_OK && vcount > 0) {
*out_vresults = vresults;
*out_vcount = vcount;
}
}
}
if (args_doc) {
yyjson_doc_free(args_doc);
}
return type_error;
return status;
}

static bool search_graph_arg_present(const char *args, const char *name) {
Expand Down Expand Up @@ -5134,9 +5148,9 @@ static char *handle_search_graph(cbm_mcp_server_t *srv, const char *args) {
cbm_vector_result_t *vresults = NULL;
int vcount = 0;
bool sq_present = false;
bool sq_type_error = run_semantic_query_core(args, store, project, semantic_materialize_limit,
&vresults, &vcount, &sq_present);
if (sq_type_error) {
sq_run_status_t sq_status = run_semantic_query_core(
args, store, project, semantic_materialize_limit, &vresults, &vcount, &sq_present);
if (sq_status != SQ_RUN_OK) {
if (fields_owner) {
yyjson_doc_free(fields_owner);
}
Expand All @@ -5146,11 +5160,17 @@ static char *handle_search_graph(cbm_mcp_server_t *srv, const char *args) {
free(qn_pattern);
free(file_pattern);
free(relationship);
if (sq_status == SQ_RUN_STORE_ERROR) {
return cbm_mcp_text_result(
"semantic search failed: the vector index could not be scanned — see the "
"server log for the SQLite error; reindex the project if it persists.",
true);
}
return cbm_mcp_text_result(
"semantic_query must be an array of keyword strings, e.g. "
"[\"send\",\"pubsub\",\"publish\"] — not a single string. Split your query "
"into individual keywords; each is scored independently via per-keyword "
"min-cosine.",
"[\"send\",\"pubsub\",\"publish\"] — not a single string, and every element "
"must be a string. Split your query into individual keywords; each is scored "
"independently via per-keyword min-cosine.",
true);
}

Expand Down
114 changes: 85 additions & 29 deletions src/store/store.c
Original file line number Diff line number Diff line change
Expand Up @@ -10018,25 +10018,42 @@ static int vs_build_keyword_vectors(cbm_store_t *s, const char *project, const c
return actual_kw;
}

/* L2 norm of every quantized keyword vector, computed once per search
* instead of once per (node, keyword) pair inside the scoring loop. */
static void vs_keyword_norms(const int8_t (*kw_vecs)[VS_VEC_DIM], int actual_kw, double *kw_norms) {
for (int k = 0; k < actual_kw; k++) {
int32_t ma = 0;
for (int d = 0; d < VS_VEC_DIM; d++) {
ma += (int32_t)kw_vecs[k][d] * (int32_t)kw_vecs[k][d];
}
kw_norms[k] = sqrt((double)ma);
}
}

/* Compute the per-keyword min cosine score between a node's int8 vector and
* each of the query vectors. Returns 0.0 if the node vector is unavailable
* or mis-sized. */
* each of the query vectors. `kw_norms[k]` is the precomputed L2 norm of
* kw_vecs[k]; the node norm is computed once per node. The integer sums are
* exact, so factoring the norms out of the pair loop leaves every score
* bit-identical to the fused form. Returns 0.0 if the node vector is
* unavailable or mis-sized. */
static double vs_min_cosine_score(const int8_t *node_vec, int node_vec_len,
const int8_t (*kw_vecs)[VS_VEC_DIM], int actual_kw) {
const int8_t (*kw_vecs)[VS_VEC_DIM], const double *kw_norms,
int actual_kw) {
if (!node_vec || node_vec_len != VS_VEC_DIM) {
return 0.0;
}
int32_t mb = 0;
for (int d = 0; d < VS_VEC_DIM; d++) {
mb += (int32_t)node_vec[d] * (int32_t)node_vec[d];
}
double node_norm = sqrt((double)mb);
double min_score = CBM_STORE_UNIT_POS_D;
for (int k = 0; k < actual_kw; k++) {
int32_t dot = 0;
int32_t ma = 0;
int32_t mb = 0;
for (int d = 0; d < VS_VEC_DIM; d++) {
dot += (int32_t)kw_vecs[k][d] * (int32_t)node_vec[d];
ma += (int32_t)kw_vecs[k][d] * (int32_t)kw_vecs[k][d];
mb += (int32_t)node_vec[d] * (int32_t)node_vec[d];
}
double denom = sqrt((double)ma) * sqrt((double)mb);
double denom = kw_norms[k] * node_norm;
double cos_k = denom > CBM_STORE_DENOM_EPS_D ? (double)dot / denom : 0.0;
if (cos_k < min_score) {
min_score = cos_k;
Expand All @@ -10045,35 +10062,69 @@ static double vs_min_cosine_score(const int8_t *node_vec, int node_vec_len,
return min_score;
}

/* Append one candidate row read from the scan statement into the result
* vector. Grows the results array geometrically on demand. Returns the
* (possibly grown) results pointer, or NULL on allocation failure. */
static cbm_vector_result_t *vs_append_result(cbm_vector_result_t *results, int *count, int *cap,
sqlite3_stmt *stmt,
const int8_t (*kw_vecs)[VS_VEC_DIM], int actual_kw) {
/* Append one candidate row read from the scan statement into `*results`,
* growing the array geometrically on demand. Returns CBM_STORE_OK, or
* CBM_STORE_ERR on allocation failure — in which case `*results` still owns
* exactly `*count` complete rows (a row is only counted once every string
* copy succeeded), so the caller's usual free path stays valid. */
static int vs_append_result(cbm_vector_result_t **results, int *count, int *cap, sqlite3_stmt *stmt,
const int8_t (*kw_vecs)[VS_VEC_DIM], const double *kw_norms,
int actual_kw) {
if (*count >= *cap) {
int nc = *cap < CBM_SZ_16 ? CBM_SZ_16 : *cap * ST_COL_2;
cbm_vector_result_t *grown = realloc(results, (size_t)nc * sizeof(cbm_vector_result_t));
cbm_vector_result_t *grown = realloc(*results, (size_t)nc * sizeof(cbm_vector_result_t));
if (!grown) {
return NULL;
return CBM_STORE_ERR;
}
results = grown;
*results = grown;
*cap = nc;
}
int idx = (*count)++;
results[idx].node_id = sqlite3_column_int64(stmt, 0);
cbm_vector_result_t *row = &(*results)[*count];
row->node_id = sqlite3_column_int64(stmt, 0);
const char *name = (const char *)sqlite3_column_text(stmt, SKIP_ONE);
const char *qn = (const char *)sqlite3_column_text(stmt, ST_COL_2);
const char *fp = (const char *)sqlite3_column_text(stmt, ST_COL_3);
const char *label = (const char *)sqlite3_column_text(stmt, ST_COL_4);
results[idx].name = name ? strdup(name) : strdup("");
results[idx].qualified_name = qn ? strdup(qn) : strdup("");
results[idx].file_path = fp ? strdup(fp) : strdup("");
results[idx].label = label ? strdup(label) : strdup("");
row->name = strdup(name ? name : "");
row->qualified_name = strdup(qn ? qn : "");
row->file_path = strdup(fp ? fp : "");
row->label = strdup(label ? label : "");
if (!row->name || !row->qualified_name || !row->file_path || !row->label) {
free(row->name);
free(row->qualified_name);
free(row->file_path);
free(row->label);
return CBM_STORE_ERR;
}
const int8_t *node_vec = (const int8_t *)sqlite3_column_blob(stmt, ST_COL_6);
int node_vec_len = sqlite3_column_bytes(stmt, ST_COL_6);
results[idx].score = vs_min_cosine_score(node_vec, node_vec_len, kw_vecs, actual_kw);
return results;
row->score = vs_min_cosine_score(node_vec, node_vec_len, kw_vecs, kw_norms, actual_kw);
(*count)++;
return CBM_STORE_OK;
}

/* A lean index carries no node_vectors table at all. That is an empty
* semantic universe, not a scan failure, and callers must be able to tell
* the two apart — so probe the schema before preparing the scan. Returns
* CBM_STORE_OK (present), CBM_STORE_NOT_FOUND (absent) or CBM_STORE_ERR. */
static int vs_probe_node_vectors_table(cbm_store_t *s) {
sqlite3_stmt *probe = NULL;
const char *sql =
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='node_vectors' LIMIT 1;";
if (sqlite3_prepare_v2(s->db, sql, SQLITE_AUTO_LEN, &probe, NULL) != SQLITE_OK) {
(void)fprintf(stderr, "vector_search: %s\n", sqlite3_errmsg(s->db));
return CBM_STORE_ERR;
}
int step_rc = sqlite3_step(probe);
sqlite3_finalize(probe);
if (step_rc == SQLITE_ROW) {
return CBM_STORE_OK;
}
if (step_rc == SQLITE_DONE) {
return CBM_STORE_NOT_FOUND;
}
(void)fprintf(stderr, "vector_search: %s\n", sqlite3_errmsg(s->db));
return CBM_STORE_ERR;
}

static int vs_ranked_result_cmp(const void *lhs, const void *rhs) {
Expand All @@ -10097,11 +10148,18 @@ int cbm_store_vector_search(cbm_store_t *s, const char *project, const char **ke
return CBM_STORE_ERR;
}

int probe_rc = vs_probe_node_vectors_table(s);
if (probe_rc != CBM_STORE_OK) {
return probe_rc;
}

int8_t kw_vecs[VS_MAX_KW][VS_VEC_DIM];
int actual_kw = vs_build_keyword_vectors(s, project, keywords, keyword_count, kw_vecs);
if (actual_kw == 0) {
return CBM_STORE_OK;
}
double kw_norms[VS_MAX_KW];
vs_keyword_norms(kw_vecs, actual_kw, kw_norms);

/* Use the first keyword for a cheap candidate ordering, then score each
* materialized candidate by min-cosine across every keyword. */
Expand Down Expand Up @@ -10141,13 +10199,11 @@ int cbm_store_vector_search(cbm_store_t *s, const char *project, const char **ke
sqlite3_bind_int(stmt, ST_COL_3, fetch_limit);
while ((step_rc = sqlite3_step(stmt)) == SQLITE_ROW) {
omitted_score_ceiling = sqlite3_column_double(stmt, ST_COL_5);
cbm_vector_result_t *grown =
vs_append_result(results, &count, &cap, stmt, kw_vecs, actual_kw);
if (!grown) {
if (vs_append_result(&results, &count, &cap, stmt, kw_vecs, kw_norms, actual_kw) !=
CBM_STORE_OK) {
step_rc = SQLITE_NOMEM;
break;
}
results = grown;
}
if (step_rc != SQLITE_DONE) {
char rc_buf[VS_STR_BUF];
Expand Down
6 changes: 5 additions & 1 deletion src/store/store.h
Original file line number Diff line number Diff line change
Expand Up @@ -1048,7 +1048,11 @@ typedef struct {
/* Search for nodes similar to the given query keywords using stored RI vectors.
* Builds a merged query vector from the keywords, then does cosine scan via
* the cbm_cosine_i8 SQL function joined with the nodes table.
* Returns results sorted by score DESC. Caller must free with cbm_store_free_vector_results. */
* Returns CBM_STORE_OK with results sorted by score DESC (possibly zero),
* CBM_STORE_NOT_FOUND when the store carries no node_vectors table (lean
* index — an empty universe, not a fault), or CBM_STORE_ERR when the scan
* itself failed; callers must not render CBM_STORE_ERR as zero matches.
* Caller must free with cbm_store_free_vector_results. */
int cbm_store_vector_search(cbm_store_t *s, const char *project, const char **keywords,
int keyword_count, int limit, cbm_vector_result_t **out,
int *out_count);
Expand Down
Loading
Loading