From 425d417d7f4919e0fc210a8ba100cbcad64397d9 Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Wed, 24 Jun 2026 12:50:17 -0600 Subject: [PATCH 1/7] feat: add hybrid search with v2 API connection and metadata filter support for chroma --- .../vector/ChromaVectorDatabaseEngine.java | 482 +++++++++++++++++- ...omaVectorQueryFilterTranslationHelper.java | 223 ++++++++ 2 files changed, 678 insertions(+), 27 deletions(-) create mode 100644 src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index d640e3fb4c1..341abe52b74 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -34,6 +34,7 @@ import java.util.ArrayList; import java.util.Date; import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Properties; @@ -54,6 +55,7 @@ import prerna.engine.api.IModelEngine; import prerna.engine.api.VectorDatabaseTypeEnum; import prerna.om.Insight; +import prerna.query.querystruct.filters.IQueryFilter; import prerna.sablecc2.om.execptions.SemossPixelException; import prerna.security.HttpHelperUtility; import prerna.util.Constants; @@ -66,17 +68,62 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { public static final String CHROMA_CLASSNAME = "CHROMA_COLLECTION_NAME"; public static final String COLLECTION_ID = "COLLECTION_ID"; + /** SMSS keys for the Chroma v2 tenant/database namespace (optional; sensible defaults applied). */ + public static final String TENANT = "TENANT"; + public static final String DB_NAME = "DB_NAME"; + + // v2 REST path fragments: {url}api/v2/tenants/{tenant}/databases/{db}/collections[/{id}{action}] + private static final String TENANTS = "api/v2/tenants"; + private static final String DATABASES = "/databases"; + private static final String COLLECTIONS = "/collections"; + private static final String DEFAULT_TENANT = "default_tenant"; + private static final String DEFAULT_DATABASE = "default_database"; + private final String API_TOKEN_KEY = "X-Chroma-Token"; - + private final String API_ADD = "/add"; private final String API_DELETE = "/delete"; private final String API_QUERY = "/query"; - + private String url = null; private String apiKey = null; + private String tenant = null; + private String dbName = null; private String className = null; private String collectionID = null; + /** SMSS key to enable hybrid (vector + keyword) search on this engine. */ + public static final String USE_HYBRID_SEARCH = "USE_HYBRID_SEARCH"; + + /** + * SMSS key for the vector similarity weight used in hybrid RRF scoring (0.0-1.0). + * The keyword weight is derived as {@code 1 - vectorWeight}. Defaults to {@code 0.5}. + */ + public static final String HYBRID_VECTOR_WEIGHT = "HYBRID_VECTOR_WEIGHT"; + + /** + * SMSS key for the minimum BM25 keyword score required for the keyword ranking to + * participate in RRF scoring. When the highest BM25 score across all candidates + * falls at or below this threshold (i.e. the query matched no candidate text), the + * keyword ranking is skipped and results are ordered by vector similarity alone. + * Defaults to {@code 0.0}. + */ + public static final String HYBRID_KEYWORD_GATE_THRESHOLD = "HYBRID_KEYWORD_GATE_THRESHOLD"; + + private static final double DEFAULT_HYBRID_VECTOR_WEIGHT = 0.5; + private static final double DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD = 0.0; + // standard RRF dampening constant + private static final int RRF_K = 60; + // BM25 tuning constants (Lucene defaults) + private static final double BM25_K1 = 1.2; + private static final double BM25_B = 0.75; + // transient per-candidate key used to carry the Chroma vector distance through ranking; stripped before return + private static final String HYBRID_DISTANCE_KEY = "_hybrid_distance"; + + private boolean useHybridSearch = false; + private double hybridVectorWeight = DEFAULT_HYBRID_VECTOR_WEIGHT; + private double hybridKeywordGateThreshold = DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD; + @Override public void open(Properties smssProp) throws Exception { super.open(smssProp); @@ -87,13 +134,71 @@ public void open(Properties smssProp) throws Exception { } this.apiKey = smssProp.getProperty(Constants.API_KEY); this.className = smssProp.getProperty(CHROMA_CLASSNAME); + this.tenant = smssProp.getProperty(TENANT); + this.dbName = smssProp.getProperty(DB_NAME); + + this.useHybridSearch = Boolean.parseBoolean(this.smssProp.getProperty(USE_HYBRID_SEARCH, "false")); + + if (this.smssProp.containsKey(HYBRID_VECTOR_WEIGHT)) { + try { + double parsedVectorWeight = Double.parseDouble(this.smssProp.getProperty(HYBRID_VECTOR_WEIGHT)); + if (parsedVectorWeight >= 0.0 && parsedVectorWeight <= 1.0) { + this.hybridVectorWeight = parsedVectorWeight; + } else { + classLogger.warn("HYBRID_VECTOR_WEIGHT '{}' must be between 0.0 and 1.0 (inclusive); defaulting to {}", + parsedVectorWeight, DEFAULT_HYBRID_VECTOR_WEIGHT); + } + } catch (NumberFormatException e) { + classLogger.warn("HYBRID_VECTOR_WEIGHT '{}' is not a valid number; defaulting to {}", + this.smssProp.getProperty(HYBRID_VECTOR_WEIGHT), DEFAULT_HYBRID_VECTOR_WEIGHT, e); + } + } + + if (this.smssProp.containsKey(HYBRID_KEYWORD_GATE_THRESHOLD)) { + try { + double parsedGateThreshold = Double.parseDouble(this.smssProp.getProperty(HYBRID_KEYWORD_GATE_THRESHOLD)); + if (parsedGateThreshold >= 0.0) { + this.hybridKeywordGateThreshold = parsedGateThreshold; + } else { + classLogger.warn("HYBRID_KEYWORD_GATE_THRESHOLD '{}' must be >= 0; defaulting to {}", + parsedGateThreshold, DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD); + } + } catch (NumberFormatException e) { + classLogger.warn("HYBRID_KEYWORD_GATE_THRESHOLD '{}' is not a valid number; defaulting to {}", + this.smssProp.getProperty(HYBRID_KEYWORD_GATE_THRESHOLD), DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD, e); + } + } // create or fetch collection Id from the Chroma DB this.collectionID = createCollection(this.className); } /** - * + * Build the Chroma v2 collections endpoint: + * {@code {url}api/v2/tenants/{tenant}/databases/{database}/collections}. Falls back to + * the default tenant/database when not configured. + */ + public static String collections(String url, String tenant, String database) { + if (tenant == null || tenant.isEmpty()) { + tenant = DEFAULT_TENANT; + } + if (database == null || database.isEmpty()) { + database = DEFAULT_DATABASE; + } + return new StringBuilder(url).append(TENANTS).append("/").append(tenant).append(DATABASES).append("/") + .append(database).append(COLLECTIONS).toString(); + } + + /** + * Build a Chroma v2 collection action endpoint, e.g. {@code .../collections/{id}/query}. + */ + public static String collection(String url, String tenant, String database, String collectionId, String action) { + return new StringBuilder(collections(url, tenant, database)).append("/").append(collectionId).append(action) + .toString(); + } + + /** + * * @param collectionName */ private String createCollection(String collectionName) { @@ -102,6 +207,7 @@ private String createCollection(String collectionName) { // if not create a collection and get the ID collectionName = collectionName.replaceAll(" ", "_"); Gson gson = new GsonBuilder().setPrettyPrinting().create(); + String collectionsUrl = collections(this.url, this.tenant, this.dbName); Map headersMap = new HashMap<>(); if (this.apiKey != null && !this.apiKey.isEmpty()) { headersMap.put(API_TOKEN_KEY, this.apiKey); @@ -109,10 +215,10 @@ private String createCollection(String collectionName) { } else { headersMap = null; } - + String nearestNeigborResponse = null; try { - nearestNeigborResponse = HttpHelperUtility.getRequest(this.url, headersMap, null, null, null); + nearestNeigborResponse = HttpHelperUtility.getRequest(collectionsUrl, headersMap, null, null, null); } catch(Exception e) { classLogger.error("Unable to create connection"); throw new SemossPixelException("Unable to create connection"); @@ -130,9 +236,9 @@ private String createCollection(String collectionName) { Map collectionNameToCreate = new HashMap<>(); collectionNameToCreate.put("name", collectionName); String body = gson.toJson(collectionNameToCreate); - nearestNeigborResponse = HttpHelperUtility.postRequestStringBody(this.url, headersMap, body, ContentType.APPLICATION_JSON, null, null, null); + nearestNeigborResponse = HttpHelperUtility.postRequestStringBody(collectionsUrl, headersMap, body, ContentType.APPLICATION_JSON, null, null, null); Map responseMap = gson.fromJson(nearestNeigborResponse, new TypeToken>() {}.getType()); - + return (String) responseMap.get("id"); } @@ -204,7 +310,8 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT headersMap = null; } - String response = HttpHelperUtility.postRequestStringBody(this.url + this.collectionID + API_ADD, + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_ADD), headersMap, body, ContentType.APPLICATION_JSON, null, null, null); List fileStatusList = new ArrayList<>(); //TODO: let us add validation by looking at the response @@ -260,9 +367,8 @@ public void removeDocument(List fileNames, Map parameter for (int fileIndex = 0; fileIndex < sourceNames.size(); fileIndex++) { String fileName = fileNames.get(fileIndex); - // Delete document in ChromaDB using their ID, but to get the ID we need to find - // the ID of a document first. Check the delete API call params - // http://localhost:5000/api/v1/collections/{}/delete + // Delete document in ChromaDB by matching the Source metadata via the v2 delete API: + // {url}api/v2/tenants/{tenant}/databases/{db}/collections/{id}/delete Map fileNamesForDelete = new HashMap<>(); Map sourceProperty = new HashMap<>(); @@ -283,7 +389,8 @@ public void removeDocument(List fileNames, Map parameter headersMap = null; } - String response = HttpHelperUtility.postRequestStringBody(this.url + this.collectionID + API_DELETE, + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_DELETE), headersMap, body, ContentType.APPLICATION_JSON, null, null, null); //TODO: let us add validation by looking at the response @@ -310,7 +417,9 @@ public void removeDocument(List fileNames, Map parameter } @Override - public List> nearestNeighborCall(Insight insight, String searchStatement, Number limit, Map parameters) { + @SuppressWarnings("unchecked") + public List> nearestNeighborCall(Insight insight, String searchStatement, Number limit, + Map parameters) { if (insight == null) { throw new IllegalArgumentException("Insight must be provided to run Model Engine Encoder"); } @@ -320,20 +429,25 @@ public List> nearestNeighborCall(Insight insight, String sea if (limit == null) { limit = 3; } - + Gson gson = new Gson(); List vector = getEmbeddingsDouble(searchStatement, insight); + Map where = buildWhere(parameters); + + if (this.useHybridSearch) { + return executeHybridRrfSearch(searchStatement, vector, limit.intValue(), where); + } + Map query = new HashMap<>(); List> queryEmbeddings = new ArrayList<>(); - // this is done to put a list of embeddings inside another list otherwise the - // API throws error. - queryEmbeddings.add(vector); - - // List> metadatas = new ArrayList<>(); add metadata filter - query.put("query_texts", searchStatement); + // nest the embedding inside a list as the API expects + queryEmbeddings.add(vector); query.put("n_results", limit); query.put("query_embeddings", queryEmbeddings); + if (where != null) { + query.put("where", where); + } String body = gson.toJson(query); Map headersMap = new HashMap<>(); @@ -343,17 +457,331 @@ public List> nearestNeighborCall(Insight insight, String sea } else { headersMap = null; } - - String nearestNeigborResponse = HttpHelperUtility.postRequestStringBody(this.url + this.collectionID + API_QUERY, + + String nearestNeigborResponse = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_QUERY), headersMap, body, ContentType.APPLICATION_JSON, null, null, null); Map responseMap = gson.fromJson(nearestNeigborResponse, new TypeToken>() {}.getType()); - - // Retrieve the metadatas list response - List> resultMap = (List>) responseMap.get("metadatas"); - return (List>) resultMap.get(0); + if (responseMap == null) { + throw new SemossPixelException("Failed to query Chroma collection."); + } + + // v2 returns parallel metadatas/distances lists; flatten into rows with Score + Distance + List> results = new ArrayList<>(); + List>> metadatas = (List>>) responseMap.get("metadatas"); + List> distances = (List>) responseMap.get("distances"); + if (metadatas != null && !metadatas.isEmpty() && metadatas.get(0) != null) { + List> metadata = metadatas.get(0); + List distance = (distances != null && !distances.isEmpty()) ? distances.get(0) : null; + for (int i = 0; i < metadata.size(); i++) { + Map row = new LinkedHashMap<>(); + if (distance != null && i < distance.size() && distance.get(i) != null) { + double d = distance.get(i); + row.put("Score", toScore(d)); + row.put("Distance", d); + } + row.putAll(metadata.get(i)); + results.add(row); + } + } + return results; } - + + /** + * Convert a Chroma distance to a higher-is-better similarity score. For cosine distance + * this is {@code 1 - distance}; for other metrics (e.g. L2) a monotonic {@code 1/(1+distance)} + * keeps higher = more similar. + * + * @param distance the raw Chroma distance + * @return a similarity score where higher is better + */ + private double toScore(double distance) { + if (this.distanceMethod != null && this.distanceMethod.toLowerCase().contains("cosine")) { + return 1.0 - distance; + } + return 1.0 / (1.0 + distance); + } + + /** + * Build a Chroma {@code where} clause from the SEMOSS {@code filters}/{@code metaFilters} + * parameters. Both lists are combined with AND. Returns {@code null} when there is nothing + * to filter on. + * + * @param parameters the reactor parameter map (may be {@code null}) + * @return a Chroma {@code where} map, or {@code null} + */ + @SuppressWarnings("unchecked") + private Map buildWhere(Map parameters) { + if (parameters == null) { + return null; + } + List allFilters = new ArrayList<>(); + Object filters = parameters.get(AbstractVectorDatabaseEngine.FILTERS_KEY); + Object metaFilters = parameters.get(AbstractVectorDatabaseEngine.METADATA_FILTERS_KEY); + if (filters instanceof List) { + allFilters.addAll((List) filters); + } + if (metaFilters instanceof List) { + allFilters.addAll((List) metaFilters); + } + return ChromaVectorQueryFilterTranslationHelper.toWhere(allFilters); + } + + /** + * Executes a hybrid vector + keyword search using Weighted Reciprocal Rank Fusion + * (RRF). Chroma has no native keyword/full-text scoring, so a candidate pool is + * fetched by vector similarity and each candidate's stored {@code Content} is scored + * with BM25 in Java. The two rankings are combined with: + * 
+	 *   rrfScore = vectorWeight / (k + vectorRank) + keywordWeight / (k + keywordRank)
+	 *
+ * where {@code k = 60}. The vector weight is configured via {@code HYBRID_VECTOR_WEIGHT} + * (default 0.5) and the keyword weight is {@code 1 - vectorWeight}. If the highest BM25 + * score across all candidates is at or below {@code HYBRID_KEYWORD_GATE_THRESHOLD}, the + * keyword ranking is skipped and results fall back to pure vector order. + *

+ * This re-ranks the vector candidate pool; a document that is a strong keyword match + * but outside the top vector candidates is not surfaced. This mirrors the vector-first + * behavior of the PGVector hybrid implementation. + *

+ * + * @param searchStatement the user's query string + * @param vector pre-computed embedding for {@code searchStatement} + * @param limit maximum number of results to return + * @param where optional Chroma {@code where} filter applied to the candidate pool; may be {@code null} + * @return results ranked by weighted RRF score, descending + */ + private List> executeHybridRrfSearch(String searchStatement, List vector, int limit, + Map where) { + int candidateLimit = Math.max(limit * 10, 100); + + List> candidates = queryChromaCandidates(vector, candidateLimit, where); + if (candidates.isEmpty()) { + return candidates; + } + int n = candidates.size(); + + // 1. BM25 keyword score per candidate, index-aligned with candidates + double[] keywordScores = bm25Scores(searchStatement, candidates); + + // 2. derive per-dimension ranks (position 0 = best): vector by ascending distance, + // keyword by descending BM25 score + List byVector = new ArrayList<>(n); + List byKeyword = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + byVector.add(i); + byKeyword.add(i); + } + byVector.sort((a, b) -> Double.compare(distanceOf(candidates.get(a)), distanceOf(candidates.get(b)))); + byKeyword.sort((a, b) -> Double.compare(keywordScores[b], keywordScores[a])); + + int[] vectorRank = new int[n]; + int[] keywordRank = new int[n]; + for (int rank = 0; rank < n; rank++) { + vectorRank[byVector.get(rank)] = rank; + keywordRank[byKeyword.get(rank)] = rank; + } + + // 3. gate: if the query matched no candidate text, skip keyword ranking + double maxKeywordScore = 0.0; + for (double score : keywordScores) { + if (score > maxKeywordScore) { + maxKeywordScore = score; + } + } + boolean useKeyword = maxKeywordScore > this.hybridKeywordGateThreshold; + + double vectorWeight = this.hybridVectorWeight; + double keywordWeight = 1.0 - vectorWeight; + classLogger.debug("Chroma hybrid RRF for query '{}': candidates={}, vectorWeight={}, keywordWeight={}, useKeyword={}, maxKeywordScore={}", + searchStatement, n, vectorWeight, keywordWeight, useKeyword, maxKeywordScore); + + // 4. weighted RRF score per candidate + double[] rrfScores = new double[n]; + for (int i = 0; i < n; i++) { + rrfScores[i] = vectorWeight / (RRF_K + vectorRank[i] + 1.0); + if (useKeyword) { + rrfScores[i] += keywordWeight / (RRF_K + keywordRank[i] + 1.0); + } + } + + // 5. return top-N by RRF score, attaching Score and stripping the transient distance + List sortedIndices = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + sortedIndices.add(i); + } + sortedIndices.sort((a, b) -> Double.compare(rrfScores[b], rrfScores[a])); + + List> results = new ArrayList<>(Math.min(limit, n)); + for (int i = 0; i < Math.min(limit, n); i++) { + int idx = sortedIndices.get(i); + Map row = candidates.get(idx); + row.put("Score", rrfScores[idx]); + row.remove(HYBRID_DISTANCE_KEY); + results.add(row); + } + return results; + } + + /** + * Fetches a candidate pool from Chroma ordered by vector similarity. Unlike the + * standard {@code nearestNeighborCall} this also captures the {@code distances} the + * Chroma response returns (carried on each row under {@link #HYBRID_DISTANCE_KEY}) so + * the candidates can be ranked by vector similarity during fusion. + * + * @param vector the query embedding + * @param nResults the candidate pool size to request + * @param where optional Chroma {@code where} filter; may be {@code null} + * @return candidate rows (metadata maps), each tagged with its vector distance; empty if none + */ + @SuppressWarnings("unchecked") + private List> queryChromaCandidates(List vector, int nResults, Map where) { + Gson gson = new Gson(); + Map query = new HashMap<>(); + List> queryEmbeddings = new ArrayList<>(); + queryEmbeddings.add(vector); + query.put("n_results", nResults); + query.put("query_embeddings", queryEmbeddings); + if (where != null) { + query.put("where", where); + } + String body = gson.toJson(query); + + Map headersMap = new HashMap<>(); + if (this.apiKey != null && !this.apiKey.isEmpty()) { + headersMap.put(API_TOKEN_KEY, this.apiKey); + headersMap.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); + } else { + headersMap = null; + } + + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_QUERY), + headersMap, body, ContentType.APPLICATION_JSON, null, null, null); + + Map responseMap = gson.fromJson(response, new TypeToken>() {}.getType()); + + List>> metadatas = (List>>) responseMap.get("metadatas"); + if (metadatas == null || metadatas.isEmpty() || metadatas.get(0) == null) { + return new ArrayList<>(); + } + List> rows = metadatas.get(0); + + // distances are returned parallel to metadatas; tag each row so vector ranking can use it + List> distances = (List>) responseMap.get("distances"); + List rowDistances = (distances != null && !distances.isEmpty()) ? distances.get(0) : null; + for (int i = 0; i < rows.size(); i++) { + double distance = (rowDistances != null && i < rowDistances.size() && rowDistances.get(i) != null) + ? rowDistances.get(i) + : Double.MAX_VALUE; + rows.get(i).put(HYBRID_DISTANCE_KEY, distance); + } + return rows; + } + + /** + * @param row a candidate row tagged by {@link #queryChromaCandidates} + * @return the vector distance for the row, or {@link Double#MAX_VALUE} if missing + */ + private double distanceOf(Map row) { + Object distance = row.get(HYBRID_DISTANCE_KEY); + return (distance instanceof Number) ? ((Number) distance).doubleValue() : Double.MAX_VALUE; + } + + /** + * Computes a BM25 keyword relevance score for each candidate against the query. The + * corpus statistics (document frequency, average document length) are derived from the + * candidate pool itself, which is sufficient for re-ranking. Returns an array of scores + * index-aligned with {@code candidates}; a candidate that matches no query term scores 0. + * + * @param query the user's query string + * @param candidates the candidate rows whose {@code Content} is scored + * @return BM25 scores aligned with {@code candidates} + */ + private double[] bm25Scores(String query, List> candidates) { + int n = candidates.size(); + double[] scores = new double[n]; + + List queryTerms = tokenize(query); + if (queryTerms.isEmpty() || n == 0) { + return scores; + } + + // term frequencies per document + document lengths + List> docTermFreqs = new ArrayList<>(n); + int[] docLengths = new int[n]; + double totalLength = 0.0; + for (int i = 0; i < n; i++) { + Object content = candidates.get(i).get(VectorDatabaseCSVTable.CONTENT); + List tokens = tokenize(content == null ? "" : content.toString()); + Map termFreq = new HashMap<>(); + for (String token : tokens) { + termFreq.put(token, termFreq.getOrDefault(token, 0) + 1); + } + docTermFreqs.add(termFreq); + docLengths[i] = tokens.size(); + totalLength += tokens.size(); + } + double avgDocLength = totalLength / n; + if (avgDocLength <= 0.0) { + return scores; + } + + // document frequency for each unique query term, over the candidate pool + Map docFreq = new HashMap<>(); + for (String term : queryTerms) { + if (docFreq.containsKey(term)) { + continue; + } + int df = 0; + for (Map termFreq : docTermFreqs) { + if (termFreq.containsKey(term)) { + df++; + } + } + docFreq.put(term, df); + } + + // BM25 score per document (only unique query terms contribute) + for (int i = 0; i < n; i++) { + Map termFreq = docTermFreqs.get(i); + double score = 0.0; + for (Map.Entry entry : docFreq.entrySet()) { + int f = termFreq.getOrDefault(entry.getKey(), 0); + if (f == 0) { + continue; + } + int df = entry.getValue(); + double idf = Math.log(1.0 + (n - df + 0.5) / (df + 0.5)); + double denom = f + BM25_K1 * (1.0 - BM25_B + BM25_B * docLengths[i] / avgDocLength); + score += idf * (f * (BM25_K1 + 1.0)) / denom; + } + scores[i] = score; + } + return scores; + } + + /** + * Lower-cases and splits text into alphanumeric tokens. No external tokenizer/stemmer + * is used so this introduces no new dependencies. + * + * @param text the text to tokenize + * @return the list of tokens (possibly empty) + */ + private List tokenize(String text) { + List tokens = new ArrayList<>(); + if (text == null || text.isEmpty()) { + return tokens; + } + for (String token : text.toLowerCase().split("[^a-z0-9]+")) { + if (!token.isEmpty()) { + tokens.add(token); + } + } + return tokens; + } + @Override public List> listDocuments(Map parameters) { //TODO: needs to grab 'Source' from the database diff --git a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java new file mode 100644 index 00000000000..39a392a017f --- /dev/null +++ b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java @@ -0,0 +1,223 @@ +/******************************************************************************* + * Copyright 2015 Defense Health Agency (DHA) + * + * If your use of this software does not include any GPLv2 components: + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * ---------------------------------------------------------------------------- + * If your use of this software includes any GPLv2 components: + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * as published by the Free Software Foundation; either version 2 + * of the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + *******************************************************************************/ +package prerna.engine.impl.vector; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import prerna.query.querystruct.filters.AbstractListFilter; +import prerna.query.querystruct.filters.IQueryFilter; +import prerna.query.querystruct.filters.SimpleQueryFilter; +import prerna.query.querystruct.filters.SimpleQueryFilter.FILTER_TYPE; +import prerna.query.querystruct.selectors.IQuerySelector; +import prerna.sablecc2.om.nounmeta.NounMetadata; + +/** + * Translates SEMOSS {@link IQueryFilter}s into a Chroma {@code where} clause. + *

+ * Chroma expects a nested JSON-style structure, which this helper produces as a + * {@link Map}: + *

+ * 
+ *   {"Source": {"$in": ["a.txt","b.txt"]}}
+ *   {"Modality": {"$eq": "text"}}
+ *   {"$and": [ {"Source": {"$eq":"a.txt"}}, {"Modality": {"$eq":"text"}} ]}
+ *   {"$or":  [ {"Modality": {"$eq":"text"}}, {"Modality": {"$eq":"image"}} ]}
+ *
+ *

+ * Supports SIMPLE column-to-values filters (mapped to {@code $eq}/{@code $in}, + * {@code $ne}/{@code $nin}, {@code $gt}/{@code $gte}/{@code $lt}/{@code $lte}) and + * arbitrarily nested {@code AND}/{@code OR} groups via recursion. Unsupported filter + * shapes are skipped rather than throwing. + *

+ */ +public final class ChromaVectorQueryFilterTranslationHelper { + + private static final String AND = "$and"; + private static final String OR = "$or"; + + private ChromaVectorQueryFilterTranslationHelper() { + } + + /** + * Combine a list of top-level filters (e.g. {@code filters} + {@code metaFilters}) into a + * single Chroma {@code where} clause, AND-ing them together. + * + * @param filters the SEMOSS filters to translate; may be {@code null}/empty + * @return the Chroma {@code where} map, or {@code null} when there is nothing to filter on + */ + public static Map toWhere(List filters) { + if (filters == null || filters.isEmpty()) { + return null; + } + List> clauses = new ArrayList<>(); + for (IQueryFilter filter : filters) { + Map clause = translate(filter); + if (clause != null) { + clauses.add(clause); + } + } + return combine(AND, clauses); + } + + private static Map translate(IQueryFilter filter) { + if (filter == null) { + return null; + } + switch (filter.getQueryFilterType()) { + case SIMPLE: + return translateSimple((SimpleQueryFilter) filter); + case AND: + return translateGroup(AND, (AbstractListFilter) filter); + case OR: + return translateGroup(OR, (AbstractListFilter) filter); + default: + return null; + } + } + + private static Map translateGroup(String operator, AbstractListFilter filter) { + List> clauses = new ArrayList<>(); + List filterList = filter.getFilterList(); + if (filterList != null) { + for (IQueryFilter child : filterList) { + Map clause = translate(child); + if (clause != null) { + clauses.add(clause); + } + } + } + return combine(operator, clauses); + } + + /** + * Collapse a set of clauses under a boolean operator. A single clause is returned as-is + * (Chroma rejects a one-element {@code $and}/{@code $or}); an empty set yields {@code null}. + */ + private static Map combine(String operator, List> clauses) { + if (clauses.isEmpty()) { + return null; + } + if (clauses.size() == 1) { + return clauses.get(0); + } + Map combined = new HashMap<>(); + combined.put(operator, clauses); + return combined; + } + + private static Map translateSimple(SimpleQueryFilter filter) { + // Only column-to-values filters (e.g. Source == "a.txt") map cleanly to a Chroma where clause + if (filter.getSimpleFilterType() != FILTER_TYPE.COL_TO_VALUES) { + return null; + } + String column = extractColumn(filter.getLComparison()); + if (column == null) { + return null; + } + List values = normalizeToList(filter.getRComparison().getValue()); + if (values.isEmpty()) { + return null; + } + Map condition = buildCondition(filter.getComparator(), values); + if (condition == null) { + return null; + } + Map clause = new HashMap<>(); + clause.put(column, condition); + return clause; + } + + private static Map buildCondition(String comparator, List values) { + boolean multi = values.size() > 1; + Map condition = new HashMap<>(); + switch (comparator) { + case "==": + case "=": + condition.put(multi ? "$in" : "$eq", multi ? values : values.get(0)); + break; + case "!=": + case "<>": + condition.put(multi ? "$nin" : "$ne", multi ? values : values.get(0)); + break; + case ">": + condition.put("$gt", values.get(0)); + break; + case ">=": + condition.put("$gte", values.get(0)); + break; + case "<": + condition.put("$lt", values.get(0)); + break; + case "<=": + condition.put("$lte", values.get(0)); + break; + default: + return null; + } + return condition; + } + + /** + * Resolve the column name from the left comparison of a simple filter. + */ + private static String extractColumn(NounMetadata leftComparison) { + if (leftComparison == null) { + return null; + } + Object value = leftComparison.getValue(); + if (value instanceof IQuerySelector) { + return ((IQuerySelector) value).getQueryStructName(); + } + return value == null ? null : value.toString(); + } + + /** + * Normalize a right-comparison value to a flat list, handling both single values and + * collections (e.g. a {@code Vector} of values for an {@code IN} filter). + */ + private static List normalizeToList(Object value) { + List values = new ArrayList<>(); + if (value == null) { + return values; + } + if (value instanceof Collection) { + for (Object element : (Collection) value) { + if (element != null) { + values.add(element); + } + } + } else { + values.add(value); + } + return values; + } +} From 8a5ddd41b647699c0a4248cc3a29be4ea8d74926 Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Thu, 25 Jun 2026 09:37:30 -0600 Subject: [PATCH 2/7] feat: dedupe chroma engine helpers and remove unchecked-cast suppressions --- .../vector/ChromaVectorDatabaseEngine.java | 286 ++++++++---------- 1 file changed, 128 insertions(+), 158 deletions(-) diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index 341abe52b74..ac6cca1aeb5 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -32,6 +32,7 @@ import java.nio.file.Paths; import java.text.SimpleDateFormat; import java.util.ArrayList; +import java.util.Comparator; import java.util.Date; import java.util.HashMap; import java.util.LinkedHashMap; @@ -47,7 +48,6 @@ import org.apache.logging.log4j.Logger; import com.google.gson.Gson; -import com.google.gson.GsonBuilder; import com.google.gson.reflect.TypeToken; import prerna.cluster.util.ClusterUtil; @@ -118,7 +118,7 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private static final double BM25_K1 = 1.2; private static final double BM25_B = 0.75; // transient per-candidate key used to carry the Chroma vector distance through ranking; stripped before return - private static final String HYBRID_DISTANCE_KEY = "_hybrid_distance"; + private static final String DISTANCE_KEY = "_distance"; private boolean useHybridSearch = false; private double hybridVectorWeight = DEFAULT_HYBRID_VECTOR_WEIGHT; @@ -197,6 +197,22 @@ public static String collection(String url, String tenant, String database, Stri .toString(); } + /** + * Build the request headers for a Chroma call. Returns {@code null} (i.e. no headers) when no + * API key is configured, which is the common case for a local/open Chroma instance. + * + * @return the header map, or {@code null} when no API key is set + */ + private Map buildHeaders() { + if (this.apiKey == null || this.apiKey.isEmpty()) { + return null; + } + Map headers = new HashMap<>(); + headers.put(API_TOKEN_KEY, this.apiKey); + headers.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); + return headers; + } + /** * * @param collectionName @@ -206,15 +222,9 @@ private String createCollection(String collectionName) { // if available, get the ID // if not create a collection and get the ID collectionName = collectionName.replaceAll(" ", "_"); - Gson gson = new GsonBuilder().setPrettyPrinting().create(); + Gson gson = new Gson(); String collectionsUrl = collections(this.url, this.tenant, this.dbName); - Map headersMap = new HashMap<>(); - if (this.apiKey != null && !this.apiKey.isEmpty()) { - headersMap.put(API_TOKEN_KEY, this.apiKey); - headersMap.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); - } else { - headersMap = null; - } + Map headersMap = buildHeaders(); String nearestNeigborResponse = null; try { @@ -302,17 +312,9 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT String body = new Gson().toJson(vectors); - Map headersMap = new HashMap<>(); - if (this.apiKey != null && !this.apiKey.isEmpty()) { - headersMap.put(API_TOKEN_KEY, this.apiKey); - headersMap.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); - } else { - headersMap = null; - } - String response = HttpHelperUtility.postRequestStringBody( collection(this.url, this.tenant, this.dbName, this.collectionID, API_ADD), - headersMap, body, ContentType.APPLICATION_JSON, null, null, null); + buildHeaders(), body, ContentType.APPLICATION_JSON, null, null, null); List fileStatusList = new ArrayList<>(); //TODO: let us add validation by looking at the response for (Map.Entry entry : fileRecordCountMap.entrySet()) { @@ -381,17 +383,9 @@ public void removeDocument(List fileNames, Map parameter String body = new Gson().toJson(fileNamesForDelete); - Map headersMap = new HashMap<>(); - if (this.apiKey != null && !this.apiKey.isEmpty()) { - headersMap.put(API_TOKEN_KEY, this.apiKey); - headersMap.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); - } else { - headersMap = null; - } - String response = HttpHelperUtility.postRequestStringBody( collection(this.url, this.tenant, this.dbName, this.collectionID, API_DELETE), - headersMap, body, ContentType.APPLICATION_JSON, null, null, null); + buildHeaders(), body, ContentType.APPLICATION_JSON, null, null, null); //TODO: let us add validation by looking at the response @@ -417,7 +411,6 @@ public void removeDocument(List fileNames, Map parameter } @Override - @SuppressWarnings("unchecked") public List> nearestNeighborCall(Insight insight, String searchStatement, Number limit, Map parameters) { if (insight == null) { @@ -430,8 +423,6 @@ public List> nearestNeighborCall(Insight insight, String sea limit = 3; } - Gson gson = new Gson(); - List vector = getEmbeddingsDouble(searchStatement, insight); Map where = buildWhere(parameters); @@ -439,51 +430,19 @@ public List> nearestNeighborCall(Insight insight, String sea return executeHybridRrfSearch(searchStatement, vector, limit.intValue(), where); } - Map query = new HashMap<>(); - List> queryEmbeddings = new ArrayList<>(); - // nest the embedding inside a list as the API expects - queryEmbeddings.add(vector); - query.put("n_results", limit); - query.put("query_embeddings", queryEmbeddings); - if (where != null) { - query.put("where", where); - } - String body = gson.toJson(query); - - Map headersMap = new HashMap<>(); - if (this.apiKey != null && !this.apiKey.isEmpty()) { - headersMap.put(API_TOKEN_KEY, this.apiKey); - headersMap.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); - } else { - headersMap = null; - } - - String nearestNeigborResponse = HttpHelperUtility.postRequestStringBody( - collection(this.url, this.tenant, this.dbName, this.collectionID, API_QUERY), - headersMap, body, ContentType.APPLICATION_JSON, null, null, null); - - Map responseMap = gson.fromJson(nearestNeigborResponse, new TypeToken>() {}.getType()); - if (responseMap == null) { - throw new SemossPixelException("Failed to query Chroma collection."); - } - - // v2 returns parallel metadatas/distances lists; flatten into rows with Score + Distance - List> results = new ArrayList<>(); - List>> metadatas = (List>>) responseMap.get("metadatas"); - List> distances = (List>) responseMap.get("distances"); - if (metadatas != null && !metadatas.isEmpty() && metadatas.get(0) != null) { - List> metadata = metadatas.get(0); - List distance = (distances != null && !distances.isEmpty()) ? distances.get(0) : null; - for (int i = 0; i < metadata.size(); i++) { - Map row = new LinkedHashMap<>(); - if (distance != null && i < distance.size() && distance.get(i) != null) { - double d = distance.get(i); - row.put("Score", toScore(d)); - row.put("Distance", d); - } - row.putAll(metadata.get(i)); - results.add(row); + // vector-only: return each candidate with its vector Score + raw Distance, metadata after + List> candidates = queryVectors(vector, limit.intValue(), where); + List> results = new ArrayList<>(candidates.size()); + for (Map candidate : candidates) { + Object distance = candidate.remove(DISTANCE_KEY); + Map row = new LinkedHashMap<>(); + if (distance instanceof Number) { + double d = ((Number) distance).doubleValue(); + row.put("Score", toScore(d)); + row.put("Distance", d); } + row.putAll(candidate); + results.add(row); } return results; } @@ -511,21 +470,32 @@ private double toScore(double distance) { * @param parameters the reactor parameter map (may be {@code null}) * @return a Chroma {@code where} map, or {@code null} */ - @SuppressWarnings("unchecked") private Map buildWhere(Map parameters) { if (parameters == null) { return null; } List allFilters = new ArrayList<>(); - Object filters = parameters.get(AbstractVectorDatabaseEngine.FILTERS_KEY); - Object metaFilters = parameters.get(AbstractVectorDatabaseEngine.METADATA_FILTERS_KEY); - if (filters instanceof List) { - allFilters.addAll((List) filters); + collectFilters(parameters.get(AbstractVectorDatabaseEngine.FILTERS_KEY), allFilters); + collectFilters(parameters.get(AbstractVectorDatabaseEngine.METADATA_FILTERS_KEY), allFilters); + return ChromaVectorQueryFilterTranslationHelper.toWhere(allFilters); + } + + /** + * Append any {@link IQueryFilter}s found in {@code value} (expected to be a list) to + * {@code target}, checking each element's type so no unchecked cast is needed. + * + * @param value the raw parameter value (may be {@code null} or a non-list) + * @param target the list to append matched filters to + */ + private void collectFilters(Object value, List target) { + if (!(value instanceof List)) { + return; } - if (metaFilters instanceof List) { - allFilters.addAll((List) metaFilters); + for (Object element : (List) value) { + if (element instanceof IQueryFilter) { + target.add((IQueryFilter) element); + } } - return ChromaVectorQueryFilterTranslationHelper.toWhere(allFilters); } /** @@ -556,40 +526,24 @@ private List> executeHybridRrfSearch(String searchStatement, Map where) { int candidateLimit = Math.max(limit * 10, 100); - List> candidates = queryChromaCandidates(vector, candidateLimit, where); + List> candidates = queryVectors(vector, candidateLimit, where); if (candidates.isEmpty()) { return candidates; } int n = candidates.size(); - // 1. BM25 keyword score per candidate, index-aligned with candidates double[] keywordScores = bm25Scores(searchStatement, candidates); - // 2. derive per-dimension ranks (position 0 = best): vector by ascending distance, - // keyword by descending BM25 score - List byVector = new ArrayList<>(n); - List byKeyword = new ArrayList<>(n); - for (int i = 0; i < n; i++) { - byVector.add(i); - byKeyword.add(i); - } - byVector.sort((a, b) -> Double.compare(distanceOf(candidates.get(a)), distanceOf(candidates.get(b)))); - byKeyword.sort((a, b) -> Double.compare(keywordScores[b], keywordScores[a])); - - int[] vectorRank = new int[n]; - int[] keywordRank = new int[n]; - for (int rank = 0; rank < n; rank++) { - vectorRank[byVector.get(rank)] = rank; - keywordRank[byKeyword.get(rank)] = rank; - } + // Rank candidates independently by vector similarity (ascending distance) and keyword score + // (descending BM25); rank 0 = best on each dimension. + List byVector = sortedIndices(n, + (a, b) -> Double.compare(distanceOf(candidates.get(a)), distanceOf(candidates.get(b)))); + List byKeyword = sortedIndices(n, (a, b) -> Double.compare(keywordScores[b], keywordScores[a])); + int[] vectorRank = toRanks(byVector); + int[] keywordRank = toRanks(byKeyword); - // 3. gate: if the query matched no candidate text, skip keyword ranking - double maxKeywordScore = 0.0; - for (double score : keywordScores) { - if (score > maxKeywordScore) { - maxKeywordScore = score; - } - } + // Drop the keyword signal when the query matched no candidate text. + double maxKeywordScore = keywordScores[byKeyword.get(0)]; boolean useKeyword = maxKeywordScore > this.hybridKeywordGateThreshold; double vectorWeight = this.hybridVectorWeight; @@ -597,7 +551,6 @@ private List> executeHybridRrfSearch(String searchStatement, classLogger.debug("Chroma hybrid RRF for query '{}': candidates={}, vectorWeight={}, keywordWeight={}, useKeyword={}, maxKeywordScore={}", searchStatement, n, vectorWeight, keywordWeight, useKeyword, maxKeywordScore); - // 4. weighted RRF score per candidate double[] rrfScores = new double[n]; for (int i = 0; i < n; i++) { rrfScores[i] = vectorWeight / (RRF_K + vectorRank[i] + 1.0); @@ -606,86 +559,93 @@ private List> executeHybridRrfSearch(String searchStatement, } } - // 5. return top-N by RRF score, attaching Score and stripping the transient distance - List sortedIndices = new ArrayList<>(n); + // Return the top-N by RRF score, attaching the fused Score and stripping the transient distance. + List ranked = sortedIndices(n, (a, b) -> Double.compare(rrfScores[b], rrfScores[a])); + int resultCount = Math.min(limit, n); + List> results = new ArrayList<>(resultCount); + for (int i = 0; i < resultCount; i++) { + Map row = candidates.get(ranked.get(i)); + row.put("Score", rrfScores[ranked.get(i)]); + row.remove(DISTANCE_KEY); + results.add(row); + } + return results; + } + + /** + * Returns the indices {@code [0, n)} sorted by {@code order} (position 0 = best). + */ + private static List sortedIndices(int n, Comparator order) { + List indices = new ArrayList<>(n); for (int i = 0; i < n; i++) { - sortedIndices.add(i); + indices.add(i); } - sortedIndices.sort((a, b) -> Double.compare(rrfScores[b], rrfScores[a])); + indices.sort(order); + return indices; + } - List> results = new ArrayList<>(Math.min(limit, n)); - for (int i = 0; i < Math.min(limit, n); i++) { - int idx = sortedIndices.get(i); - Map row = candidates.get(idx); - row.put("Score", rrfScores[idx]); - row.remove(HYBRID_DISTANCE_KEY); - results.add(row); + /** + * Inverts a best-first index ordering into a per-index rank array (rank 0 = best). + */ + private static int[] toRanks(List ordered) { + int[] rank = new int[ordered.size()]; + for (int position = 0; position < ordered.size(); position++) { + rank[ordered.get(position)] = position; } - return results; + return rank; } /** - * Fetches a candidate pool from Chroma ordered by vector similarity. Unlike the - * standard {@code nearestNeighborCall} this also captures the {@code distances} the - * Chroma response returns (carried on each row under {@link #HYBRID_DISTANCE_KEY}) so - * the candidates can be ranked by vector similarity during fusion. + * Runs a vector query against the collection and returns the first query's candidate rows. + * Each row is the chunk's metadata map, additionally tagged with its raw vector distance under + * {@link #DISTANCE_KEY} (read back via {@link #distanceOf}). Shared by the vector-only and + * hybrid search paths. * - * @param vector the query embedding - * @param nResults the candidate pool size to request - * @param where optional Chroma {@code where} filter; may be {@code null} - * @return candidate rows (metadata maps), each tagged with its vector distance; empty if none + * @param vector the query embedding + * @param nResults the number of results to request + * @param where optional Chroma {@code where} filter; may be {@code null} + * @return candidate rows, each tagged with its vector distance; empty if the query had no hits */ - @SuppressWarnings("unchecked") - private List> queryChromaCandidates(List vector, int nResults, Map where) { + private List> queryVectors(List vector, int nResults, Map where) { Gson gson = new Gson(); Map query = new HashMap<>(); List> queryEmbeddings = new ArrayList<>(); + // nest the embedding inside a list as the API expects queryEmbeddings.add(vector); query.put("n_results", nResults); query.put("query_embeddings", queryEmbeddings); if (where != null) { query.put("where", where); } - String body = gson.toJson(query); - - Map headersMap = new HashMap<>(); - if (this.apiKey != null && !this.apiKey.isEmpty()) { - headersMap.put(API_TOKEN_KEY, this.apiKey); - headersMap.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); - } else { - headersMap = null; - } - String response = HttpHelperUtility.postRequestStringBody( + String responseBody = HttpHelperUtility.postRequestStringBody( collection(this.url, this.tenant, this.dbName, this.collectionID, API_QUERY), - headersMap, body, ContentType.APPLICATION_JSON, null, null, null); - - Map responseMap = gson.fromJson(response, new TypeToken>() {}.getType()); + buildHeaders(), gson.toJson(query), ContentType.APPLICATION_JSON, null, null, null); - List>> metadatas = (List>>) responseMap.get("metadatas"); - if (metadatas == null || metadatas.isEmpty() || metadatas.get(0) == null) { + ChromaQueryResponse parsed = gson.fromJson(responseBody, ChromaQueryResponse.class); + if (parsed == null) { + throw new SemossPixelException("Failed to query Chroma collection."); + } + if (parsed.metadatas == null || parsed.metadatas.isEmpty() || parsed.metadatas.get(0) == null) { return new ArrayList<>(); } - List> rows = metadatas.get(0); - // distances are returned parallel to metadatas; tag each row so vector ranking can use it - List> distances = (List>) responseMap.get("distances"); - List rowDistances = (distances != null && !distances.isEmpty()) ? distances.get(0) : null; + List> rows = parsed.metadatas.get(0); + List distances = (parsed.distances != null && !parsed.distances.isEmpty()) ? parsed.distances.get(0) + : null; for (int i = 0; i < rows.size(); i++) { - double distance = (rowDistances != null && i < rowDistances.size() && rowDistances.get(i) != null) - ? rowDistances.get(i) - : Double.MAX_VALUE; - rows.get(i).put(HYBRID_DISTANCE_KEY, distance); + Double distance = (distances != null && i < distances.size()) ? distances.get(i) : null; + rows.get(i).put(DISTANCE_KEY, distance); } return rows; } /** - * @param row a candidate row tagged by {@link #queryChromaCandidates} - * @return the vector distance for the row, or {@link Double#MAX_VALUE} if missing + * @param row a candidate row tagged by {@link #queryVectors} + * @return the vector distance for the row, or {@link Double#MAX_VALUE} if missing (sorts last) */ private double distanceOf(Map row) { - Object distance = row.get(HYBRID_DISTANCE_KEY); + Object distance = row.get(DISTANCE_KEY); return (distance instanceof Number) ? ((Number) distance).doubleValue() : Double.MAX_VALUE; } @@ -830,4 +790,14 @@ public VectorDatabaseTypeEnum getVectorDatabaseType() { return VectorDatabaseTypeEnum.CHROMA; } + /** + * Typed view of a Chroma query response. Gson resolves the nested generics from these + * declared fields, so the response can be deserialized without unchecked casts. Both + * {@code metadatas} and {@code distances} are lists-of-lists (one inner list per query). + */ + private static class ChromaQueryResponse { + private List>> metadatas; + private List> distances; + } + } \ No newline at end of file From a1087371fa37444215ee2f8cc0bf2b2c4220824c Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Thu, 25 Jun 2026 11:24:37 -0600 Subject: [PATCH 3/7] refactor: simplify chroma engine internals and validate delete responses --- .../vector/ChromaVectorDatabaseEngine.java | 182 +++++------------- ...omaVectorQueryFilterTranslationHelper.java | 46 +---- 2 files changed, 59 insertions(+), 169 deletions(-) diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index ac6cca1aeb5..f9b35f248ad 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -92,32 +92,20 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private String className = null; private String collectionID = null; - /** SMSS key to enable hybrid (vector + keyword) search on this engine. */ + /** SMSS key: enable hybrid (vector + keyword) search. */ public static final String USE_HYBRID_SEARCH = "USE_HYBRID_SEARCH"; - - /** - * SMSS key for the vector similarity weight used in hybrid RRF scoring (0.0-1.0). - * The keyword weight is derived as {@code 1 - vectorWeight}. Defaults to {@code 0.5}. - */ + /** SMSS key: vector weight in RRF (0.0-1.0); keyword weight is {@code 1 - this}. Default 0.5. */ public static final String HYBRID_VECTOR_WEIGHT = "HYBRID_VECTOR_WEIGHT"; - - /** - * SMSS key for the minimum BM25 keyword score required for the keyword ranking to - * participate in RRF scoring. When the highest BM25 score across all candidates - * falls at or below this threshold (i.e. the query matched no candidate text), the - * keyword ranking is skipped and results are ordered by vector similarity alone. - * Defaults to {@code 0.0}. - */ + /** SMSS key: min BM25 score for keyword ranking to apply; below it, results are vector-only. Default 0.0. */ public static final String HYBRID_KEYWORD_GATE_THRESHOLD = "HYBRID_KEYWORD_GATE_THRESHOLD"; private static final double DEFAULT_HYBRID_VECTOR_WEIGHT = 0.5; private static final double DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD = 0.0; - // standard RRF dampening constant private static final int RRF_K = 60; // BM25 tuning constants (Lucene defaults) private static final double BM25_K1 = 1.2; private static final double BM25_B = 0.75; - // transient per-candidate key used to carry the Chroma vector distance through ranking; stripped before return + // transient key holding each candidate's vector distance during ranking; stripped before return private static final String DISTANCE_KEY = "_distance"; private boolean useHybridSearch = false; @@ -173,11 +161,7 @@ public void open(Properties smssProp) throws Exception { this.collectionID = createCollection(this.className); } - /** - * Build the Chroma v2 collections endpoint: - * {@code {url}api/v2/tenants/{tenant}/databases/{database}/collections}. Falls back to - * the default tenant/database when not configured. - */ + /** Chroma v2 collections endpoint: {@code {url}api/v2/tenants/{tenant}/databases/{db}/collections}. */ public static String collections(String url, String tenant, String database) { if (tenant == null || tenant.isEmpty()) { tenant = DEFAULT_TENANT; @@ -189,20 +173,13 @@ public static String collections(String url, String tenant, String database) { .append(database).append(COLLECTIONS).toString(); } - /** - * Build a Chroma v2 collection action endpoint, e.g. {@code .../collections/{id}/query}. - */ + /** Chroma v2 collection action endpoint, e.g. {@code .../collections/{id}/query}. */ public static String collection(String url, String tenant, String database, String collectionId, String action) { return new StringBuilder(collections(url, tenant, database)).append("/").append(collectionId).append(action) .toString(); } - /** - * Build the request headers for a Chroma call. Returns {@code null} (i.e. no headers) when no - * API key is configured, which is the common case for a local/open Chroma instance. - * - * @return the header map, or {@code null} when no API key is set - */ + /** Request headers for a Chroma call, or {@code null} when no API key is configured. */ private Map buildHeaders() { if (this.apiKey == null || this.apiKey.isEmpty()) { return null; @@ -213,14 +190,8 @@ private Map buildHeaders() { return headers; } - /** - * - * @param collectionName - */ + /** Return the id of the named collection, creating it if it does not exist. */ private String createCollection(String collectionName) { - // check to see if the collection is available - // if available, get the ID - // if not create a collection and get the ID collectionName = collectionName.replaceAll(" ", "_"); Gson gson = new Gson(); String collectionsUrl = collections(this.url, this.tenant, this.dbName); @@ -369,9 +340,7 @@ public void removeDocument(List fileNames, Map parameter for (int fileIndex = 0; fileIndex < sourceNames.size(); fileIndex++) { String fileName = fileNames.get(fileIndex); - // Delete document in ChromaDB by matching the Source metadata via the v2 delete API: - // {url}api/v2/tenants/{tenant}/databases/{db}/collections/{id}/delete - + // delete chunks matching the Source metadata (v2 delete API) Map fileNamesForDelete = new HashMap<>(); Map sourceProperty = new HashMap<>(); @@ -381,14 +350,25 @@ public void removeDocument(List fileNames, Map parameter fileNamesForDelete.put("where", sourceProperty); - String body = new Gson().toJson(fileNamesForDelete); + Gson gson = new Gson(); + String body = gson.toJson(fileNamesForDelete); String response = HttpHelperUtility.postRequestStringBody( collection(this.url, this.tenant, this.dbName, this.collectionID, API_DELETE), buildHeaders(), body, ContentType.APPLICATION_JSON, null, null, null); - //TODO: let us add validation by looking at the response - + // Chroma returns {"deleted": N}; log it and warn when nothing matched + String source = sourceProperty.get("Source"); + ChromaDeleteResponse deleteResponse = gson.fromJson(response, ChromaDeleteResponse.class); + int deleted = (deleteResponse != null && deleteResponse.deleted != null) ? deleteResponse.deleted : 0; + if (deleted > 0) { + classLogger.info("Removed {} record(s) from Chroma collection '{}' for source '{}'", deleted, + this.className, source); + } else { + classLogger.warn("No records matched source '{}' in Chroma collection '{}' during delete", source, + this.className); + } + String documentName = Paths.get(fileName).getFileName().toString(); // remove the physical documents File documentFile = new File(this.schemaFolder.getAbsolutePath() + FILE_SEPARATOR + indexClass + FILE_SEPARATOR + "documents", documentName); @@ -447,14 +427,7 @@ public List> nearestNeighborCall(Insight insight, String sea return results; } - /** - * Convert a Chroma distance to a higher-is-better similarity score. For cosine distance - * this is {@code 1 - distance}; for other metrics (e.g. L2) a monotonic {@code 1/(1+distance)} - * keeps higher = more similar. - * - * @param distance the raw Chroma distance - * @return a similarity score where higher is better - */ + /** Higher-is-better similarity score from a Chroma distance: {@code 1 - d} for cosine, else {@code 1/(1+d)}. */ private double toScore(double distance) { if (this.distanceMethod != null && this.distanceMethod.toLowerCase().contains("cosine")) { return 1.0 - distance; @@ -462,14 +435,7 @@ private double toScore(double distance) { return 1.0 / (1.0 + distance); } - /** - * Build a Chroma {@code where} clause from the SEMOSS {@code filters}/{@code metaFilters} - * parameters. Both lists are combined with AND. Returns {@code null} when there is nothing - * to filter on. - * - * @param parameters the reactor parameter map (may be {@code null}) - * @return a Chroma {@code where} map, or {@code null} - */ + /** Build a Chroma {@code where} clause from the {@code filters}/{@code metaFilters} params (AND-combined), or {@code null}. */ private Map buildWhere(Map parameters) { if (parameters == null) { return null; @@ -480,13 +446,7 @@ private Map buildWhere(Map parameters) { return ChromaVectorQueryFilterTranslationHelper.toWhere(allFilters); } - /** - * Append any {@link IQueryFilter}s found in {@code value} (expected to be a list) to - * {@code target}, checking each element's type so no unchecked cast is needed. - * - * @param value the raw parameter value (may be {@code null} or a non-list) - * @param target the list to append matched filters to - */ + /** Append any {@link IQueryFilter}s in {@code value} (if it is a list) to {@code target}. */ private void collectFilters(Object value, List target) { if (!(value instanceof List)) { return; @@ -499,28 +459,10 @@ private void collectFilters(Object value, List target) { } /** - * Executes a hybrid vector + keyword search using Weighted Reciprocal Rank Fusion - * (RRF). Chroma has no native keyword/full-text scoring, so a candidate pool is - * fetched by vector similarity and each candidate's stored {@code Content} is scored - * with BM25 in Java. The two rankings are combined with: - * 
-	 *   rrfScore = vectorWeight / (k + vectorRank) + keywordWeight / (k + keywordRank)
-	 *
- * where {@code k = 60}. The vector weight is configured via {@code HYBRID_VECTOR_WEIGHT} - * (default 0.5) and the keyword weight is {@code 1 - vectorWeight}. If the highest BM25 - * score across all candidates is at or below {@code HYBRID_KEYWORD_GATE_THRESHOLD}, the - * keyword ranking is skipped and results fall back to pure vector order. - *

- * This re-ranks the vector candidate pool; a document that is a strong keyword match - * but outside the top vector candidates is not surfaced. This mirrors the vector-first - * behavior of the PGVector hybrid implementation. - *

- * - * @param searchStatement the user's query string - * @param vector pre-computed embedding for {@code searchStatement} - * @param limit maximum number of results to return - * @param where optional Chroma {@code where} filter applied to the candidate pool; may be {@code null} - * @return results ranked by weighted RRF score, descending + * Re-rank a vector candidate pool by fusing vector similarity with an in-app BM25 keyword score + * via weighted RRF: {@code score = vw/(k+vRank) + kw/(k+kRank)}, k=60. Keyword ranking is skipped + * when no candidate matches the query. Vector-first: a keyword-only match outside the pool is not + * surfaced (same trade-off as the PGVector hybrid). */ private List> executeHybridRrfSearch(String searchStatement, List vector, int limit, Map where) { @@ -534,15 +476,14 @@ private List> executeHybridRrfSearch(String searchStatement, double[] keywordScores = bm25Scores(searchStatement, candidates); - // Rank candidates independently by vector similarity (ascending distance) and keyword score - // (descending BM25); rank 0 = best on each dimension. + // rank by vector distance (asc) and keyword score (desc); rank 0 = best List byVector = sortedIndices(n, (a, b) -> Double.compare(distanceOf(candidates.get(a)), distanceOf(candidates.get(b)))); List byKeyword = sortedIndices(n, (a, b) -> Double.compare(keywordScores[b], keywordScores[a])); int[] vectorRank = toRanks(byVector); int[] keywordRank = toRanks(byKeyword); - // Drop the keyword signal when the query matched no candidate text. + // skip keyword signal when nothing matched the query text double maxKeywordScore = keywordScores[byKeyword.get(0)]; boolean useKeyword = maxKeywordScore > this.hybridKeywordGateThreshold; @@ -559,7 +500,7 @@ private List> executeHybridRrfSearch(String searchStatement, } } - // Return the top-N by RRF score, attaching the fused Score and stripping the transient distance. + // top-N by fused RRF score List ranked = sortedIndices(n, (a, b) -> Double.compare(rrfScores[b], rrfScores[a])); int resultCount = Math.min(limit, n); List> results = new ArrayList<>(resultCount); @@ -572,9 +513,7 @@ private List> executeHybridRrfSearch(String searchStatement, return results; } - /** - * Returns the indices {@code [0, n)} sorted by {@code order} (position 0 = best). - */ + /** Indices {@code [0, n)} sorted by {@code order} (position 0 = best). */ private static List sortedIndices(int n, Comparator order) { List indices = new ArrayList<>(n); for (int i = 0; i < n; i++) { @@ -584,9 +523,7 @@ private static List sortedIndices(int n, Comparator order) { return indices; } - /** - * Inverts a best-first index ordering into a per-index rank array (rank 0 = best). - */ + /** Invert a best-first index ordering into a per-index rank array (rank 0 = best). */ private static int[] toRanks(List ordered) { int[] rank = new int[ordered.size()]; for (int position = 0; position < ordered.size(); position++) { @@ -596,15 +533,8 @@ private static int[] toRanks(List ordered) { } /** - * Runs a vector query against the collection and returns the first query's candidate rows. - * Each row is the chunk's metadata map, additionally tagged with its raw vector distance under - * {@link #DISTANCE_KEY} (read back via {@link #distanceOf}). Shared by the vector-only and - * hybrid search paths. - * - * @param vector the query embedding - * @param nResults the number of results to request - * @param where optional Chroma {@code where} filter; may be {@code null} - * @return candidate rows, each tagged with its vector distance; empty if the query had no hits + * Run a vector query and return the first result set's rows, each tagged with its vector distance + * under {@link #DISTANCE_KEY}. Shared by the vector-only and hybrid paths. Empty if no hits. */ private List> queryVectors(List vector, int nResults, Map where) { Gson gson = new Gson(); @@ -640,24 +570,15 @@ private List> queryVectors(List vector, int nResults return rows; } - /** - * @param row a candidate row tagged by {@link #queryVectors} - * @return the vector distance for the row, or {@link Double#MAX_VALUE} if missing (sorts last) - */ + /** Vector distance tagged on a candidate row, or {@link Double#MAX_VALUE} if missing (sorts last). */ private double distanceOf(Map row) { Object distance = row.get(DISTANCE_KEY); return (distance instanceof Number) ? ((Number) distance).doubleValue() : Double.MAX_VALUE; } /** - * Computes a BM25 keyword relevance score for each candidate against the query. The - * corpus statistics (document frequency, average document length) are derived from the - * candidate pool itself, which is sufficient for re-ranking. Returns an array of scores - * index-aligned with {@code candidates}; a candidate that matches no query term scores 0. - * - * @param query the user's query string - * @param candidates the candidate rows whose {@code Content} is scored - * @return BM25 scores aligned with {@code candidates} + * BM25 keyword score per candidate (index-aligned), using corpus statistics from the candidate + * pool itself. A candidate matching no query term scores 0. */ private double[] bm25Scores(String query, List> candidates) { int n = candidates.size(); @@ -688,7 +609,7 @@ private double[] bm25Scores(String query, List> candidates) return scores; } - // document frequency for each unique query term, over the candidate pool + // document frequency per query term Map docFreq = new HashMap<>(); for (String term : queryTerms) { if (docFreq.containsKey(term)) { @@ -703,7 +624,7 @@ private double[] bm25Scores(String query, List> candidates) docFreq.put(term, df); } - // BM25 score per document (only unique query terms contribute) + // BM25 score per document for (int i = 0; i < n; i++) { Map termFreq = docTermFreqs.get(i); double score = 0.0; @@ -722,13 +643,7 @@ private double[] bm25Scores(String query, List> candidates) return scores; } - /** - * Lower-cases and splits text into alphanumeric tokens. No external tokenizer/stemmer - * is used so this introduces no new dependencies. - * - * @param text the text to tokenize - * @return the list of tokens (possibly empty) - */ + /** Lower-case and split into alphanumeric tokens (no external tokenizer/stemmer). */ private List tokenize(String text) { List tokens = new ArrayList<>(); if (text == null || text.isEmpty()) { @@ -790,14 +705,15 @@ public VectorDatabaseTypeEnum getVectorDatabaseType() { return VectorDatabaseTypeEnum.CHROMA; } - /** - * Typed view of a Chroma query response. Gson resolves the nested generics from these - * declared fields, so the response can be deserialized without unchecked casts. Both - * {@code metadatas} and {@code distances} are lists-of-lists (one inner list per query). - */ + /** Typed Chroma query response — lets Gson resolve the nested generics without unchecked casts. */ private static class ChromaQueryResponse { private List>> metadatas; private List> distances; } + /** Typed Chroma delete response: {@code {"deleted": N}}. */ + private static class ChromaDeleteResponse { + private Integer deleted; + } + } \ No newline at end of file diff --git a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java index 39a392a017f..3688cc63ce5 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java +++ b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java @@ -41,23 +41,11 @@ import prerna.sablecc2.om.nounmeta.NounMetadata; /** - * Translates SEMOSS {@link IQueryFilter}s into a Chroma {@code where} clause. - *

- * Chroma expects a nested JSON-style structure, which this helper produces as a - * {@link Map}: - *

- * 
- *   {"Source": {"$in": ["a.txt","b.txt"]}}
- *   {"Modality": {"$eq": "text"}}
- *   {"$and": [ {"Source": {"$eq":"a.txt"}}, {"Modality": {"$eq":"text"}} ]}
- *   {"$or":  [ {"Modality": {"$eq":"text"}}, {"Modality": {"$eq":"image"}} ]}
- *
- *

- * Supports SIMPLE column-to-values filters (mapped to {@code $eq}/{@code $in}, - * {@code $ne}/{@code $nin}, {@code $gt}/{@code $gte}/{@code $lt}/{@code $lte}) and - * arbitrarily nested {@code AND}/{@code OR} groups via recursion. Unsupported filter - * shapes are skipped rather than throwing. - *

+ * Translates SEMOSS {@link IQueryFilter}s into a Chroma {@code where} clause — e.g. + * {@code Filter(Source == ["a","b"])} becomes {@code {"Source": {"$in": ["a","b"]}}}, with nested + * {@code $and}/{@code $or} groups. Supports column-to-values filters + * ({@code $eq}/{@code $in}/{@code $ne}/{@code $nin}/{@code $gt}/{@code $gte}/{@code $lt}/{@code $lte}); + * unsupported filter shapes are skipped rather than throwing. */ public final class ChromaVectorQueryFilterTranslationHelper { @@ -67,13 +55,7 @@ public final class ChromaVectorQueryFilterTranslationHelper { private ChromaVectorQueryFilterTranslationHelper() { } - /** - * Combine a list of top-level filters (e.g. {@code filters} + {@code metaFilters}) into a - * single Chroma {@code where} clause, AND-ing them together. - * - * @param filters the SEMOSS filters to translate; may be {@code null}/empty - * @return the Chroma {@code where} map, or {@code null} when there is nothing to filter on - */ + /** Combine top-level filters into a single AND-ed Chroma {@code where} clause, or {@code null} if none. */ public static Map toWhere(List filters) { if (filters == null || filters.isEmpty()) { return null; @@ -118,10 +100,7 @@ private static Map translateGroup(String operator, AbstractListF return combine(operator, clauses); } - /** - * Collapse a set of clauses under a boolean operator. A single clause is returned as-is - * (Chroma rejects a one-element {@code $and}/{@code $or}); an empty set yields {@code null}. - */ + /** Collapse clauses under a boolean operator; one clause is returned as-is, none yields {@code null}. */ private static Map combine(String operator, List> clauses) { if (clauses.isEmpty()) { return null; @@ -135,7 +114,7 @@ private static Map combine(String operator, List translateSimple(SimpleQueryFilter filter) { - // Only column-to-values filters (e.g. Source == "a.txt") map cleanly to a Chroma where clause + // only column-to-values filters map to a Chroma where clause if (filter.getSimpleFilterType() != FILTER_TYPE.COL_TO_VALUES) { return null; } @@ -186,9 +165,7 @@ private static Map buildCondition(String comparator, List normalizeToList(Object value) { List values = new ArrayList<>(); if (value == null) { From 5a16a51914a23965308cd3251ed665ba8980911e Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Fri, 26 Jun 2026 09:23:39 -0600 Subject: [PATCH 4/7] feat: add persistent bm25 index to chroma for full-corpus hybrid search --- .../engine/impl/vector/ChromaBm25Index.java | 262 ++++++++++++ .../vector/ChromaVectorDatabaseEngine.java | 374 ++++++++++-------- ...omaVectorQueryFilterTranslationHelper.java | 102 +++++ 3 files changed, 581 insertions(+), 157 deletions(-) create mode 100644 src/prerna/engine/impl/vector/ChromaBm25Index.java diff --git a/src/prerna/engine/impl/vector/ChromaBm25Index.java b/src/prerna/engine/impl/vector/ChromaBm25Index.java new file mode 100644 index 00000000000..4c26422851b --- /dev/null +++ b/src/prerna/engine/impl/vector/ChromaBm25Index.java @@ -0,0 +1,262 @@ +/******************************************************************************* + * Copyright 2015 Defense Health Agency (DHA) + * + * If your use of this software does not include any GPLv2 components: + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * ---------------------------------------------------------------------------- + * If your use of this software includes any GPLv2 components: + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * as published by the Free Software Foundation; either version 2 + * of the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + *******************************************************************************/ +package prerna.engine.impl.vector; + +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +import org.apache.commons.io.FileUtils; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import com.google.gson.Gson; +import com.google.gson.reflect.TypeToken; + +/** + * Self-contained, persistent BM25 keyword index for a Chroma collection. Stores each chunk's id, + * source, metadata, and per-term frequencies so keyword retrieval runs over the full corpus — not + * just a vector candidate pool — and survives restarts. Chroma OSS has no usable native keyword + * search over REST, so this provides it app-side; thread-safe via a read/write lock. + */ +public class ChromaBm25Index { + + private static final Logger classLogger = LogManager.getLogger(ChromaBm25Index.class); + private static final Gson GSON = new Gson(); + + private static final double BM25_K1 = 1.2; + private static final double BM25_B = 0.75; + private static final String INDEX_FILE_NAME = "bm25_index.json"; + + /** Key under which {@link #search} tags each result row with its chunk id (for fusion/dedup). */ + public static final String ID_KEY = "_bm25_id"; + + /** One indexed chunk: enough to score it and return it without re-fetching from Chroma. */ + private static class Record { + private String id; + private String source; + private Map metadata; + private Map termFreqs; + private int length; + } + + private final File indexFile; + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private Map records = new LinkedHashMap<>(); + private Map docFreq = new LinkedHashMap<>(); + private double avgDocLength = 0.0; + private boolean statsDirty = true; + + public ChromaBm25Index(File dir) { + this.indexFile = new File(dir, INDEX_FILE_NAME); + } + + /** Absolute path of the backing file (for cluster push/pull). */ + public String getFilePath() { + return indexFile.getAbsolutePath(); + } + + /** Load the index from disk if a file exists; a no-op otherwise. */ + public void load() { + lock.writeLock().lock(); + try { + if (!indexFile.exists()) { + return; + } + String json = FileUtils.readFileToString(indexFile, StandardCharsets.UTF_8); + Map loaded = GSON.fromJson(json, new TypeToken>() {}.getType()); + this.records = (loaded != null) ? loaded : new LinkedHashMap<>(); + this.statsDirty = true; + } catch (Exception e) { + classLogger.error("Failed to load BM25 index from {}; starting empty", indexFile.getAbsolutePath(), e); + this.records = new LinkedHashMap<>(); + } finally { + lock.writeLock().unlock(); + } + } + + /** Persist the index to disk, creating the parent directory if needed. */ + public void save() { + lock.readLock().lock(); + try { + FileUtils.forceMkdirParent(indexFile); + FileUtils.writeStringToFile(indexFile, GSON.toJson(this.records), StandardCharsets.UTF_8); + } catch (Exception e) { + classLogger.error("Failed to save BM25 index to {}", indexFile.getAbsolutePath(), e); + } finally { + lock.readLock().unlock(); + } + } + + public boolean isEmpty() { + lock.readLock().lock(); + try { + return records.isEmpty(); + } finally { + lock.readLock().unlock(); + } + } + + /** Add (or replace) one chunk. {@code metadata} is copied so later mutations don't leak in. */ + public void addRecord(String id, String source, String content, Map metadata) { + Record record = new Record(); + record.id = id; + record.source = source; + record.metadata = (metadata != null) ? new LinkedHashMap<>(metadata) : new LinkedHashMap<>(); + record.termFreqs = new LinkedHashMap<>(); + List tokens = tokenize(content); + for (String token : tokens) { + record.termFreqs.merge(token, 1, Integer::sum); + } + record.length = tokens.size(); + + lock.writeLock().lock(); + try { + records.put(id, record); + statsDirty = true; + } finally { + lock.writeLock().unlock(); + } + } + + /** + * Remove every chunk whose source matches. + * + * @return the number of chunks removed + */ + public int removeBySource(String source) { + lock.writeLock().lock(); + try { + int before = records.size(); + records.values().removeIf(r -> source.equals(r.source)); + int removed = before - records.size(); + if (removed > 0) { + statsDirty = true; + } + return removed; + } finally { + lock.writeLock().unlock(); + } + } + + /** + * Score the query against the full corpus and return the top-{@code topK} chunks, best first. + * Each result is a copy of the chunk's metadata plus its BM25 {@code Score} and {@link #ID_KEY}; + * only chunks matching at least one query term are returned. + */ + public List> search(String query, int topK) { + List queryTerms = uniqueTokens(query); + List> results = new ArrayList<>(); + if (queryTerms.isEmpty()) { + return results; + } + + lock.readLock().lock(); + try { + refreshStatsIfNeeded(); + if (records.isEmpty() || avgDocLength <= 0.0) { + return results; + } + int n = records.size(); + for (Record record : records.values()) { + double score = scoreRecord(record, queryTerms, n); + if (score <= 0.0) { + continue; + } + Map row = new LinkedHashMap<>(record.metadata); + row.put(ID_KEY, record.id); + row.put("Score", score); + results.add(row); + } + results.sort((a, b) -> Double.compare((double) b.get("Score"), (double) a.get("Score"))); + return results.subList(0, Math.min(topK, results.size())); + } finally { + lock.readLock().unlock(); + } + } + + private double scoreRecord(Record record, List queryTerms, int n) { + double score = 0.0; + for (String term : queryTerms) { + Integer f = record.termFreqs.get(term); + if (f == null) { + continue; + } + int df = docFreq.getOrDefault(term, 0); + double idf = Math.log(1.0 + (n - df + 0.5) / (df + 0.5)); + double denom = f + BM25_K1 * (1.0 - BM25_B + BM25_B * record.length / avgDocLength); + score += idf * (f * (BM25_K1 + 1.0)) / denom; + } + return score; + } + + /** Recompute document-frequency and average length. Caller must hold a lock. */ + private void refreshStatsIfNeeded() { + if (!statsDirty) { + return; + } + docFreq = new LinkedHashMap<>(); + long totalLength = 0; + for (Record record : records.values()) { + totalLength += record.length; + for (String term : record.termFreqs.keySet()) { + docFreq.merge(term, 1, Integer::sum); + } + } + avgDocLength = records.isEmpty() ? 0.0 : (double) totalLength / records.size(); + statsDirty = false; + } + + /** Lower-case and split into alphanumeric tokens (no external tokenizer/stemmer). */ + public static List tokenize(String text) { + List tokens = new ArrayList<>(); + if (text == null || text.isEmpty()) { + return tokens; + } + for (String token : text.toLowerCase().split("[^a-z0-9]+")) { + if (!token.isEmpty()) { + tokens.add(token); + } + } + return tokens; + } + + private static List uniqueTokens(String text) { + List unique = new ArrayList<>(); + for (String token : tokenize(text)) { + if (!unique.contains(token)) { + unique.add(token); + } + } + return unique; + } +} diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index f9b35f248ad..eea1ee3a28d 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -32,7 +32,6 @@ import java.nio.file.Paths; import java.text.SimpleDateFormat; import java.util.ArrayList; -import java.util.Comparator; import java.util.Date; import java.util.HashMap; import java.util.LinkedHashMap; @@ -51,6 +50,7 @@ import com.google.gson.reflect.TypeToken; import prerna.cluster.util.ClusterUtil; +import prerna.cluster.util.CopyFilesToEngineRunner; import prerna.cluster.util.DeleteFilesFromEngineRunner; import prerna.engine.api.IModelEngine; import prerna.engine.api.VectorDatabaseTypeEnum; @@ -84,6 +84,7 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private final String API_ADD = "/add"; private final String API_DELETE = "/delete"; private final String API_QUERY = "/query"; + private final String API_GET = "/get"; private String url = null; private String apiKey = null; @@ -102,15 +103,14 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private static final double DEFAULT_HYBRID_VECTOR_WEIGHT = 0.5; private static final double DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD = 0.0; private static final int RRF_K = 60; - // BM25 tuning constants (Lucene defaults) - private static final double BM25_K1 = 1.2; - private static final double BM25_B = 0.75; // transient key holding each candidate's vector distance during ranking; stripped before return private static final String DISTANCE_KEY = "_distance"; private boolean useHybridSearch = false; private double hybridVectorWeight = DEFAULT_HYBRID_VECTOR_WEIGHT; private double hybridKeywordGateThreshold = DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD; + // persistent keyword index; populated only when hybrid search is enabled + private ChromaBm25Index bm25Index; @Override public void open(Properties smssProp) throws Exception { @@ -127,38 +127,101 @@ public void open(Properties smssProp) throws Exception { this.useHybridSearch = Boolean.parseBoolean(this.smssProp.getProperty(USE_HYBRID_SEARCH, "false")); - if (this.smssProp.containsKey(HYBRID_VECTOR_WEIGHT)) { - try { - double parsedVectorWeight = Double.parseDouble(this.smssProp.getProperty(HYBRID_VECTOR_WEIGHT)); - if (parsedVectorWeight >= 0.0 && parsedVectorWeight <= 1.0) { - this.hybridVectorWeight = parsedVectorWeight; - } else { - classLogger.warn("HYBRID_VECTOR_WEIGHT '{}' must be between 0.0 and 1.0 (inclusive); defaulting to {}", - parsedVectorWeight, DEFAULT_HYBRID_VECTOR_WEIGHT); + // weight/gate tuning only matters when hybrid search is enabled + if (this.useHybridSearch) { + if (this.smssProp.containsKey(HYBRID_VECTOR_WEIGHT)) { + try { + double parsedVectorWeight = Double.parseDouble(this.smssProp.getProperty(HYBRID_VECTOR_WEIGHT)); + if (parsedVectorWeight >= 0.0 && parsedVectorWeight <= 1.0) { + this.hybridVectorWeight = parsedVectorWeight; + } else { + classLogger.warn("HYBRID_VECTOR_WEIGHT '{}' must be between 0.0 and 1.0 (inclusive); defaulting to {}", + parsedVectorWeight, DEFAULT_HYBRID_VECTOR_WEIGHT); + } + } catch (NumberFormatException e) { + classLogger.warn("HYBRID_VECTOR_WEIGHT '{}' is not a valid number; defaulting to {}", + this.smssProp.getProperty(HYBRID_VECTOR_WEIGHT), DEFAULT_HYBRID_VECTOR_WEIGHT, e); } - } catch (NumberFormatException e) { - classLogger.warn("HYBRID_VECTOR_WEIGHT '{}' is not a valid number; defaulting to {}", - this.smssProp.getProperty(HYBRID_VECTOR_WEIGHT), DEFAULT_HYBRID_VECTOR_WEIGHT, e); } - } - if (this.smssProp.containsKey(HYBRID_KEYWORD_GATE_THRESHOLD)) { - try { - double parsedGateThreshold = Double.parseDouble(this.smssProp.getProperty(HYBRID_KEYWORD_GATE_THRESHOLD)); - if (parsedGateThreshold >= 0.0) { - this.hybridKeywordGateThreshold = parsedGateThreshold; - } else { - classLogger.warn("HYBRID_KEYWORD_GATE_THRESHOLD '{}' must be >= 0; defaulting to {}", - parsedGateThreshold, DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD); + if (this.smssProp.containsKey(HYBRID_KEYWORD_GATE_THRESHOLD)) { + try { + double parsedGateThreshold = Double.parseDouble(this.smssProp.getProperty(HYBRID_KEYWORD_GATE_THRESHOLD)); + if (parsedGateThreshold >= 0.0) { + this.hybridKeywordGateThreshold = parsedGateThreshold; + } else { + classLogger.warn("HYBRID_KEYWORD_GATE_THRESHOLD '{}' must be >= 0; defaulting to {}", + parsedGateThreshold, DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD); + } + } catch (NumberFormatException e) { + classLogger.warn("HYBRID_KEYWORD_GATE_THRESHOLD '{}' is not a valid number; defaulting to {}", + this.smssProp.getProperty(HYBRID_KEYWORD_GATE_THRESHOLD), DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD, e); } - } catch (NumberFormatException e) { - classLogger.warn("HYBRID_KEYWORD_GATE_THRESHOLD '{}' is not a valid number; defaulting to {}", - this.smssProp.getProperty(HYBRID_KEYWORD_GATE_THRESHOLD), DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD, e); } } // create or fetch collection Id from the Chroma DB this.collectionID = createCollection(this.className); + + if (this.useHybridSearch) { + initBm25Index(); + } + } + + /** + * Load the persistent BM25 index from disk, or backfill it from the existing collection if the + * index file is missing (e.g. first run after enabling hybrid, or a fresh cluster node). Failures + * are non-fatal: hybrid search simply contributes no keyword signal until the index is populated. + */ + private void initBm25Index() { + File bm25Dir = new File(this.schemaFolder.getAbsolutePath() + FILE_SEPARATOR + this.defaultIndexClass + + FILE_SEPARATOR + "bm25"); + this.bm25Index = new ChromaBm25Index(bm25Dir); + try { + this.bm25Index.load(); + if (this.bm25Index.isEmpty()) { + backfillBm25IndexFromCollection(); + } + } catch (Exception e) { + classLogger.error("Failed to initialize BM25 index for engine '{}'; keyword scoring disabled until rebuilt", + this.className, e); + } + } + + /** Rebuild the BM25 index by reading every chunk currently stored in the Chroma collection. */ + private void backfillBm25IndexFromCollection() { + Gson gson = new Gson(); + Map getBody = new HashMap<>(); + List include = new ArrayList<>(); + include.add("metadatas"); + getBody.put("include", include); + + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_GET), + buildHeaders(), gson.toJson(getBody), ContentType.APPLICATION_JSON, null, null, null); + + ChromaGetResponse parsed = gson.fromJson(response, ChromaGetResponse.class); + if (parsed == null || parsed.ids == null || parsed.metadatas == null) { + return; + } + int count = Math.min(parsed.ids.size(), parsed.metadatas.size()); + for (int i = 0; i < count; i++) { + Map metadata = parsed.metadatas.get(i); + Object content = metadata.get(VectorDatabaseCSVTable.CONTENT); + Object source = metadata.get(VectorDatabaseCSVTable.SOURCE); + this.bm25Index.addRecord(parsed.ids.get(i), source == null ? null : source.toString(), + content == null ? "" : content.toString(), metadata); + } + this.bm25Index.save(); + classLogger.info("Backfilled BM25 index for engine '{}' with {} chunk(s)", this.className, count); + } + + /** Push the BM25 index file to cloud storage when running in a cluster (mirrors document sync). */ + private void pushBm25IndexIfClustered() { + if (ClusterUtil.IS_CLUSTER && this.bm25Index != null) { + Thread.ofVirtual().start(new CopyFilesToEngineRunner(this.engineId, this.getCatalogType(), + new String[] { this.bm25Index.getFilePath() })); + } } /** Chroma v2 collections endpoint: {@code {url}api/v2/tenants/{tenant}/databases/{db}/collections}. */ @@ -313,6 +376,19 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT } + // keep the BM25 index in sync with the successful add + if (this.useHybridSearch && this.bm25Index != null && response != null && !response.trim().isEmpty()) { + for (int i = 0; i < ids.size(); i++) { + Map metadata = metadatas.get(i); + Object source = metadata.get(VectorDatabaseCSVTable.SOURCE); + Object content = metadata.get(VectorDatabaseCSVTable.CONTENT); + this.bm25Index.addRecord(ids.get(i), source == null ? null : source.toString(), + content == null ? "" : content.toString(), metadata); + } + this.bm25Index.save(); + pushBm25IndexIfClustered(); + } + return fileStatusList; } @@ -335,6 +411,7 @@ public void removeDocument(List fileNames, Map parameter } List filesToRemoveFromCloud = new ArrayList(); + boolean bm25Changed = false; // need to get the source names and then delete it based on the names for (int fileIndex = 0; fileIndex < sourceNames.size(); fileIndex++) { @@ -369,6 +446,11 @@ public void removeDocument(List fileNames, Map parameter this.className); } + // prune the same source from the BM25 index + if (this.useHybridSearch && this.bm25Index != null && this.bm25Index.removeBySource(source) > 0) { + bm25Changed = true; + } + String documentName = Paths.get(fileName).getFileName().toString(); // remove the physical documents File documentFile = new File(this.schemaFolder.getAbsolutePath() + FILE_SEPARATOR + indexClass + FILE_SEPARATOR + "documents", documentName); @@ -383,6 +465,11 @@ public void removeDocument(List fileNames, Map parameter } + if (bm25Changed) { + this.bm25Index.save(); + pushBm25IndexIfClustered(); + } + if (ClusterUtil.IS_CLUSTER) { Thread deleteFilesFromCloudThread = new Thread(new DeleteFilesFromEngineRunner(engineId, this.getCatalogType(), filesToRemoveFromCloud.stream().toArray(String[]::new))); @@ -404,7 +491,7 @@ public List> nearestNeighborCall(Insight insight, String sea } List vector = getEmbeddingsDouble(searchStatement, insight); - Map where = buildWhere(parameters); + Map where = buildWhereClause(parameters); if (this.useHybridSearch) { return executeHybridRrfSearch(searchStatement, vector, limit.intValue(), where); @@ -415,6 +502,7 @@ public List> nearestNeighborCall(Insight insight, String sea List> results = new ArrayList<>(candidates.size()); for (Map candidate : candidates) { Object distance = candidate.remove(DISTANCE_KEY); + candidate.remove(ChromaBm25Index.ID_KEY); Map row = new LinkedHashMap<>(); if (distance instanceof Number) { double d = ((Number) distance).doubleValue(); @@ -436,7 +524,7 @@ private double toScore(double distance) { } /** Build a Chroma {@code where} clause from the {@code filters}/{@code metaFilters} params (AND-combined), or {@code null}. */ - private Map buildWhere(Map parameters) { + private Map buildWhereClause(Map parameters) { if (parameters == null) { return null; } @@ -459,77 +547,119 @@ private void collectFilters(Object value, List target) { } /** - * Re-rank a vector candidate pool by fusing vector similarity with an in-app BM25 keyword score - * via weighted RRF: {@code score = vw/(k+vRank) + kw/(k+kRank)}, k=60. Keyword ranking is skipped - * when no candidate matches the query. Vector-first: a keyword-only match outside the pool is not - * surfaced (same trade-off as the PGVector hybrid). + * Hybrid search via full-corpus union: take the vector candidate pool and the BM25 keyword hits + * (from the persistent index), union them by chunk id, rank each dimension independently, and fuse + * with weighted RRF ({@code score = vw/(k+vRank) + kw/(k+kRank)}, k=60). A document contributes a + * term only for the dimension(s) it appears in, so a keyword-only match outside the vector pool can + * still surface. Keyword ranking is skipped when nothing matched the query text. */ private List> executeHybridRrfSearch(String searchStatement, List vector, int limit, Map where) { int candidateLimit = Math.max(limit * 10, 100); - List> candidates = queryVectors(vector, candidateLimit, where); - if (candidates.isEmpty()) { - return candidates; + List> vectorHits = queryVectors(vector, candidateLimit, where); + List> keywordHits = (this.bm25Index != null) + ? this.bm25Index.search(searchStatement, candidateLimit) + : new ArrayList<>(); + + // union vector + keyword hits by chunk id + Map> rowById = new LinkedHashMap<>(); + Map distanceById = new LinkedHashMap<>(); + Map keywordById = new LinkedHashMap<>(); + for (Map hit : vectorHits) { + String id = idOf(hit); + if (id == null) { + continue; + } + rowById.put(id, hit); + distanceById.put(id, distanceOf(hit)); + } + for (Map hit : keywordHits) { + // the vector side is filtered by Chroma; apply the same filter to the keyword side + if (where != null && !ChromaVectorQueryFilterTranslationHelper.matches(hit, where)) { + continue; + } + String id = idOf(hit); + if (id == null) { + continue; + } + keywordById.put(id, ((Number) hit.get("Score")).doubleValue()); + rowById.putIfAbsent(id, hit); + } + if (rowById.isEmpty()) { + return new ArrayList<>(); } - int n = candidates.size(); - - double[] keywordScores = bm25Scores(searchStatement, candidates); - - // rank by vector distance (asc) and keyword score (desc); rank 0 = best - List byVector = sortedIndices(n, - (a, b) -> Double.compare(distanceOf(candidates.get(a)), distanceOf(candidates.get(b)))); - List byKeyword = sortedIndices(n, (a, b) -> Double.compare(keywordScores[b], keywordScores[a])); - int[] vectorRank = toRanks(byVector); - int[] keywordRank = toRanks(byKeyword); - // skip keyword signal when nothing matched the query text - double maxKeywordScore = keywordScores[byKeyword.get(0)]; - boolean useKeyword = maxKeywordScore > this.hybridKeywordGateThreshold; + // independent ranks (rank 0 = best): vector by ascending distance, keyword by descending BM25 + Map vectorRank = ranksOf(distanceById, true); + boolean useKeyword = !keywordById.isEmpty() && maxValue(keywordById) > this.hybridKeywordGateThreshold; + Map keywordRank = ranksOf(keywordById, false); double vectorWeight = this.hybridVectorWeight; double keywordWeight = 1.0 - vectorWeight; - classLogger.debug("Chroma hybrid RRF for query '{}': candidates={}, vectorWeight={}, keywordWeight={}, useKeyword={}, maxKeywordScore={}", - searchStatement, n, vectorWeight, keywordWeight, useKeyword, maxKeywordScore); - - double[] rrfScores = new double[n]; - for (int i = 0; i < n; i++) { - rrfScores[i] = vectorWeight / (RRF_K + vectorRank[i] + 1.0); - if (useKeyword) { - rrfScores[i] += keywordWeight / (RRF_K + keywordRank[i] + 1.0); + classLogger.debug("Chroma hybrid RRF for query '{}': vectorHits={}, keywordHits={}, union={}, vectorWeight={}, keywordWeight={}, useKeyword={}", + searchStatement, vectorHits.size(), keywordHits.size(), rowById.size(), vectorWeight, keywordWeight, useKeyword); + + // RRF over the union: each doc contributes per dimension only where it ranks + Map rrfScores = new LinkedHashMap<>(); + for (String id : rowById.keySet()) { + double score = 0.0; + Integer vRank = vectorRank.get(id); + if (vRank != null) { + score += vectorWeight / (RRF_K + vRank + 1.0); + } + Integer kRank = keywordRank.get(id); + if (useKeyword && kRank != null) { + score += keywordWeight / (RRF_K + kRank + 1.0); } + rrfScores.put(id, score); } - // top-N by fused RRF score - List ranked = sortedIndices(n, (a, b) -> Double.compare(rrfScores[b], rrfScores[a])); - int resultCount = Math.min(limit, n); + // top-N by fused score + List rankedIds = new ArrayList<>(rrfScores.keySet()); + rankedIds.sort((a, b) -> Double.compare(rrfScores.get(b), rrfScores.get(a))); + int resultCount = Math.min(limit, rankedIds.size()); List> results = new ArrayList<>(resultCount); for (int i = 0; i < resultCount; i++) { - Map row = candidates.get(ranked.get(i)); - row.put("Score", rrfScores[ranked.get(i)]); + String id = rankedIds.get(i); + Map row = rowById.get(id); + row.put("Score", rrfScores.get(id)); row.remove(DISTANCE_KEY); + row.remove(ChromaBm25Index.ID_KEY); results.add(row); } return results; } - /** Indices {@code [0, n)} sorted by {@code order} (position 0 = best). */ - private static List sortedIndices(int n, Comparator order) { - List indices = new ArrayList<>(n); - for (int i = 0; i < n; i++) { - indices.add(i); - } - indices.sort(order); - return indices; + /** The chunk id tagged on a row (vector hits via {@link #queryVectors}, keyword hits via the index). */ + private static String idOf(Map row) { + Object id = row.get(ChromaBm25Index.ID_KEY); + return (id == null) ? null : id.toString(); + } + + /** + * Rank ids by their score (rank 0 = best). {@code ascending} ranks smaller values first (vector + * distance); {@code !ascending} ranks larger values first (keyword score). + */ + private static Map ranksOf(Map scoreById, boolean ascending) { + List order = new ArrayList<>(scoreById.keySet()); + order.sort((a, b) -> ascending ? Double.compare(scoreById.get(a), scoreById.get(b)) + : Double.compare(scoreById.get(b), scoreById.get(a))); + Map ranks = new LinkedHashMap<>(); + for (int i = 0; i < order.size(); i++) { + ranks.put(order.get(i), i); + } + return ranks; } - /** Invert a best-first index ordering into a per-index rank array (rank 0 = best). */ - private static int[] toRanks(List ordered) { - int[] rank = new int[ordered.size()]; - for (int position = 0; position < ordered.size(); position++) { - rank[ordered.get(position)] = position; + private static double maxValue(Map values) { + double max = 0.0; + for (double value : values.values()) { + if (value > max) { + max = value; + } } - return rank; + return max; } /** @@ -563,9 +693,13 @@ private List> queryVectors(List vector, int nResults List> rows = parsed.metadatas.get(0); List distances = (parsed.distances != null && !parsed.distances.isEmpty()) ? parsed.distances.get(0) : null; + List ids = (parsed.ids != null && !parsed.ids.isEmpty()) ? parsed.ids.get(0) : null; for (int i = 0; i < rows.size(); i++) { Double distance = (distances != null && i < distances.size()) ? distances.get(i) : null; rows.get(i).put(DISTANCE_KEY, distance); + if (ids != null && i < ids.size()) { + rows.get(i).put(ChromaBm25Index.ID_KEY, ids.get(i)); + } } return rows; } @@ -576,87 +710,6 @@ private double distanceOf(Map row) { return (distance instanceof Number) ? ((Number) distance).doubleValue() : Double.MAX_VALUE; } - /** - * BM25 keyword score per candidate (index-aligned), using corpus statistics from the candidate - * pool itself. A candidate matching no query term scores 0. - */ - private double[] bm25Scores(String query, List> candidates) { - int n = candidates.size(); - double[] scores = new double[n]; - - List queryTerms = tokenize(query); - if (queryTerms.isEmpty() || n == 0) { - return scores; - } - - // term frequencies per document + document lengths - List> docTermFreqs = new ArrayList<>(n); - int[] docLengths = new int[n]; - double totalLength = 0.0; - for (int i = 0; i < n; i++) { - Object content = candidates.get(i).get(VectorDatabaseCSVTable.CONTENT); - List tokens = tokenize(content == null ? "" : content.toString()); - Map termFreq = new HashMap<>(); - for (String token : tokens) { - termFreq.put(token, termFreq.getOrDefault(token, 0) + 1); - } - docTermFreqs.add(termFreq); - docLengths[i] = tokens.size(); - totalLength += tokens.size(); - } - double avgDocLength = totalLength / n; - if (avgDocLength <= 0.0) { - return scores; - } - - // document frequency per query term - Map docFreq = new HashMap<>(); - for (String term : queryTerms) { - if (docFreq.containsKey(term)) { - continue; - } - int df = 0; - for (Map termFreq : docTermFreqs) { - if (termFreq.containsKey(term)) { - df++; - } - } - docFreq.put(term, df); - } - - // BM25 score per document - for (int i = 0; i < n; i++) { - Map termFreq = docTermFreqs.get(i); - double score = 0.0; - for (Map.Entry entry : docFreq.entrySet()) { - int f = termFreq.getOrDefault(entry.getKey(), 0); - if (f == 0) { - continue; - } - int df = entry.getValue(); - double idf = Math.log(1.0 + (n - df + 0.5) / (df + 0.5)); - double denom = f + BM25_K1 * (1.0 - BM25_B + BM25_B * docLengths[i] / avgDocLength); - score += idf * (f * (BM25_K1 + 1.0)) / denom; - } - scores[i] = score; - } - return scores; - } - - /** Lower-case and split into alphanumeric tokens (no external tokenizer/stemmer). */ - private List tokenize(String text) { - List tokens = new ArrayList<>(); - if (text == null || text.isEmpty()) { - return tokens; - } - for (String token : text.toLowerCase().split("[^a-z0-9]+")) { - if (!token.isEmpty()) { - tokens.add(token); - } - } - return tokens; - } - @Override public List> listDocuments(Map parameters) { //TODO: needs to grab 'Source' from the database @@ -707,6 +760,7 @@ public VectorDatabaseTypeEnum getVectorDatabaseType() { /** Typed Chroma query response — lets Gson resolve the nested generics without unchecked casts. */ private static class ChromaQueryResponse { + private List> ids; private List>> metadatas; private List> distances; } @@ -716,4 +770,10 @@ private static class ChromaDeleteResponse { private Integer deleted; } + /** Typed Chroma {@code /get} response (flat ids + metadata, used to backfill the BM25 index). */ + private static class ChromaGetResponse { + private List ids; + private List> metadatas; + } + } \ No newline at end of file diff --git a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java index 3688cc63ce5..c4aed225fef 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java +++ b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java @@ -194,4 +194,106 @@ private static List normalizeToList(Object value) { } return values; } + + /** + * Evaluate a Chroma {@code where} clause (as produced by {@link #toWhere}) against a row's + * metadata. Used to apply the same filter to app-side BM25 results that Chroma applies to the + * vector query. Returns {@code true} when {@code where} is null/empty. + */ + public static boolean matches(Map metadata, Map where) { + return matchesNode(metadata, where); + } + + private static boolean matchesNode(Map metadata, Map node) { + if (node == null || node.isEmpty()) { + return true; + } + for (Map.Entry entry : node.entrySet()) { + String key = String.valueOf(entry.getKey()); + Object value = entry.getValue(); + if (AND.equals(key)) { + for (Object sub : asList(value)) { + if (sub instanceof Map && !matchesNode(metadata, (Map) sub)) { + return false; + } + } + } else if (OR.equals(key)) { + boolean any = false; + for (Object sub : asList(value)) { + if (sub instanceof Map && matchesNode(metadata, (Map) sub)) { + any = true; + break; + } + } + if (!any) { + return false; + } + } else if (value instanceof Map) { + if (!matchesField(metadata == null ? null : metadata.get(key), (Map) value)) { + return false; + } + } + } + return true; + } + + private static boolean matchesField(Object actual, Map condition) { + for (Map.Entry entry : condition.entrySet()) { + String op = String.valueOf(entry.getKey()); + Object operand = entry.getValue(); + switch (op) { + case "$eq": if (!valuesEqual(actual, operand)) return false; break; + case "$ne": if (valuesEqual(actual, operand)) return false; break; + case "$in": if (!containsValue(asList(operand), actual)) return false; break; + case "$nin": if (containsValue(asList(operand), actual)) return false; break; + case "$gt": case "$gte": case "$lt": case "$lte": + if (!matchesRange(op, actual, operand)) return false; + break; + default: return false; + } + } + return true; + } + + /** Numeric range comparison; fails closed (no match) when either side is not numeric. */ + private static boolean matchesRange(String op, Object actual, Object operand) { + Double a = toDouble(actual); + Double b = toDouble(operand); + if (a == null || b == null) { + return false; + } + int cmp = Double.compare(a, b); + switch (op) { + case "$gt": return cmp > 0; + case "$gte": return cmp >= 0; + case "$lt": return cmp < 0; + case "$lte": return cmp <= 0; + default: return false; + } + } + + private static Double toDouble(Object value) { + try { + return Double.parseDouble(String.valueOf(value)); + } catch (NumberFormatException e) { + return null; + } + } + + private static List asList(Object value) { + return (value instanceof List) ? (List) value : new ArrayList<>(); + } + + private static boolean valuesEqual(Object a, Object b) { + return String.valueOf(a).equals(String.valueOf(b)); + } + + private static boolean containsValue(List values, Object actual) { + for (Object value : values) { + if (valuesEqual(actual, value)) { + return true; + } + } + return false; + } } From e6901fe47564b26ee7fefbcfb70c89561c7b00dd Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Fri, 26 Jun 2026 10:20:11 -0600 Subject: [PATCH 5/7] feat: add in-memory bm25 index for full-corpus hybrid search in chroma --- .../engine/impl/vector/ChromaBm25Index.java | 75 ++------------- .../vector/ChromaVectorDatabaseEngine.java | 93 +++++++------------ 2 files changed, 41 insertions(+), 127 deletions(-) diff --git a/src/prerna/engine/impl/vector/ChromaBm25Index.java b/src/prerna/engine/impl/vector/ChromaBm25Index.java index 4c26422851b..3465daa2c3e 100644 --- a/src/prerna/engine/impl/vector/ChromaBm25Index.java +++ b/src/prerna/engine/impl/vector/ChromaBm25Index.java @@ -27,96 +27,40 @@ *******************************************************************************/ package prerna.engine.impl.vector; -import java.io.File; -import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.concurrent.locks.ReentrantReadWriteLock; -import org.apache.commons.io.FileUtils; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; - -import com.google.gson.Gson; -import com.google.gson.reflect.TypeToken; - /** - * Self-contained, persistent BM25 keyword index for a Chroma collection. Stores each chunk's id, - * source, metadata, and per-term frequencies so keyword retrieval runs over the full corpus — not - * just a vector candidate pool — and survives restarts. Chroma OSS has no usable native keyword - * search over REST, so this provides it app-side; thread-safe via a read/write lock. + * In-memory BM25 keyword index for a Chroma collection. Chroma OSS has no usable native keyword + * search over REST, so this provides it app-side. The index is derived from the chunk {@code Content} + * that already lives in Chroma — it holds no separate persisted state and is rebuilt from the + * collection on engine open. Thread-safe via a read/write lock. */ public class ChromaBm25Index { - private static final Logger classLogger = LogManager.getLogger(ChromaBm25Index.class); - private static final Gson GSON = new Gson(); - private static final double BM25_K1 = 1.2; private static final double BM25_B = 0.75; - private static final String INDEX_FILE_NAME = "bm25_index.json"; /** Key under which {@link #search} tags each result row with its chunk id (for fusion/dedup). */ public static final String ID_KEY = "_bm25_id"; /** One indexed chunk: enough to score it and return it without re-fetching from Chroma. */ private static class Record { - private String id; private String source; private Map metadata; private Map termFreqs; private int length; } - private final File indexFile; private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); - private Map records = new LinkedHashMap<>(); + private final Map records = new LinkedHashMap<>(); private Map docFreq = new LinkedHashMap<>(); private double avgDocLength = 0.0; private boolean statsDirty = true; - public ChromaBm25Index(File dir) { - this.indexFile = new File(dir, INDEX_FILE_NAME); - } - - /** Absolute path of the backing file (for cluster push/pull). */ - public String getFilePath() { - return indexFile.getAbsolutePath(); - } - - /** Load the index from disk if a file exists; a no-op otherwise. */ - public void load() { - lock.writeLock().lock(); - try { - if (!indexFile.exists()) { - return; - } - String json = FileUtils.readFileToString(indexFile, StandardCharsets.UTF_8); - Map loaded = GSON.fromJson(json, new TypeToken>() {}.getType()); - this.records = (loaded != null) ? loaded : new LinkedHashMap<>(); - this.statsDirty = true; - } catch (Exception e) { - classLogger.error("Failed to load BM25 index from {}; starting empty", indexFile.getAbsolutePath(), e); - this.records = new LinkedHashMap<>(); - } finally { - lock.writeLock().unlock(); - } - } - - /** Persist the index to disk, creating the parent directory if needed. */ - public void save() { - lock.readLock().lock(); - try { - FileUtils.forceMkdirParent(indexFile); - FileUtils.writeStringToFile(indexFile, GSON.toJson(this.records), StandardCharsets.UTF_8); - } catch (Exception e) { - classLogger.error("Failed to save BM25 index to {}", indexFile.getAbsolutePath(), e); - } finally { - lock.readLock().unlock(); - } - } - public boolean isEmpty() { lock.readLock().lock(); try { @@ -129,7 +73,6 @@ public boolean isEmpty() { /** Add (or replace) one chunk. {@code metadata} is copied so later mutations don't leak in. */ public void addRecord(String id, String source, String content, Map metadata) { Record record = new Record(); - record.id = id; record.source = source; record.metadata = (metadata != null) ? new LinkedHashMap<>(metadata) : new LinkedHashMap<>(); record.termFreqs = new LinkedHashMap<>(); @@ -187,13 +130,13 @@ public List> search(String query, int topK) { return results; } int n = records.size(); - for (Record record : records.values()) { - double score = scoreRecord(record, queryTerms, n); + for (Map.Entry entry : records.entrySet()) { + double score = scoreRecord(entry.getValue(), queryTerms, n); if (score <= 0.0) { continue; } - Map row = new LinkedHashMap<>(record.metadata); - row.put(ID_KEY, record.id); + Map row = new LinkedHashMap<>(entry.getValue().metadata); + row.put(ID_KEY, entry.getKey()); row.put("Score", score); results.add(row); } diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index eea1ee3a28d..ee81272cd2c 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -50,7 +50,6 @@ import com.google.gson.reflect.TypeToken; import prerna.cluster.util.ClusterUtil; -import prerna.cluster.util.CopyFilesToEngineRunner; import prerna.cluster.util.DeleteFilesFromEngineRunner; import prerna.engine.api.IModelEngine; import prerna.engine.api.VectorDatabaseTypeEnum; @@ -109,7 +108,7 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private boolean useHybridSearch = false; private double hybridVectorWeight = DEFAULT_HYBRID_VECTOR_WEIGHT; private double hybridKeywordGateThreshold = DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD; - // persistent keyword index; populated only when hybrid search is enabled + // in-memory keyword index (rebuilt from Chroma on open); used only when hybrid search is enabled private ChromaBm25Index bm25Index; @Override @@ -169,58 +168,38 @@ public void open(Properties smssProp) throws Exception { } /** - * Load the persistent BM25 index from disk, or backfill it from the existing collection if the - * index file is missing (e.g. first run after enabling hybrid, or a fresh cluster node). Failures - * are non-fatal: hybrid search simply contributes no keyword signal until the index is populated. + * Build the in-memory BM25 index from the chunks already stored in the Chroma collection. The + * index holds no separate persisted state — Chroma is the source of truth — so it is rebuilt on + * every engine open. Failures are non-fatal: hybrid search simply contributes no keyword signal. */ private void initBm25Index() { - File bm25Dir = new File(this.schemaFolder.getAbsolutePath() + FILE_SEPARATOR + this.defaultIndexClass - + FILE_SEPARATOR + "bm25"); - this.bm25Index = new ChromaBm25Index(bm25Dir); + this.bm25Index = new ChromaBm25Index(); try { - this.bm25Index.load(); - if (this.bm25Index.isEmpty()) { - backfillBm25IndexFromCollection(); - } - } catch (Exception e) { - classLogger.error("Failed to initialize BM25 index for engine '{}'; keyword scoring disabled until rebuilt", - this.className, e); - } - } - - /** Rebuild the BM25 index by reading every chunk currently stored in the Chroma collection. */ - private void backfillBm25IndexFromCollection() { - Gson gson = new Gson(); - Map getBody = new HashMap<>(); - List include = new ArrayList<>(); - include.add("metadatas"); - getBody.put("include", include); - - String response = HttpHelperUtility.postRequestStringBody( - collection(this.url, this.tenant, this.dbName, this.collectionID, API_GET), - buildHeaders(), gson.toJson(getBody), ContentType.APPLICATION_JSON, null, null, null); + Gson gson = new Gson(); + Map getBody = new HashMap<>(); + List include = new ArrayList<>(); + include.add("metadatas"); + getBody.put("include", include); - ChromaGetResponse parsed = gson.fromJson(response, ChromaGetResponse.class); - if (parsed == null || parsed.ids == null || parsed.metadatas == null) { - return; - } - int count = Math.min(parsed.ids.size(), parsed.metadatas.size()); - for (int i = 0; i < count; i++) { - Map metadata = parsed.metadatas.get(i); - Object content = metadata.get(VectorDatabaseCSVTable.CONTENT); - Object source = metadata.get(VectorDatabaseCSVTable.SOURCE); - this.bm25Index.addRecord(parsed.ids.get(i), source == null ? null : source.toString(), - content == null ? "" : content.toString(), metadata); - } - this.bm25Index.save(); - classLogger.info("Backfilled BM25 index for engine '{}' with {} chunk(s)", this.className, count); - } + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_GET), + buildHeaders(), gson.toJson(getBody), ContentType.APPLICATION_JSON, null, null, null); - /** Push the BM25 index file to cloud storage when running in a cluster (mirrors document sync). */ - private void pushBm25IndexIfClustered() { - if (ClusterUtil.IS_CLUSTER && this.bm25Index != null) { - Thread.ofVirtual().start(new CopyFilesToEngineRunner(this.engineId, this.getCatalogType(), - new String[] { this.bm25Index.getFilePath() })); + ChromaGetResponse parsed = gson.fromJson(response, ChromaGetResponse.class); + if (parsed == null || parsed.ids == null || parsed.metadatas == null) { + return; + } + int count = Math.min(parsed.ids.size(), parsed.metadatas.size()); + for (int i = 0; i < count; i++) { + Map metadata = parsed.metadatas.get(i); + Object content = metadata.get(VectorDatabaseCSVTable.CONTENT); + Object source = metadata.get(VectorDatabaseCSVTable.SOURCE); + this.bm25Index.addRecord(parsed.ids.get(i), source == null ? null : source.toString(), + content == null ? "" : content.toString(), metadata); + } + classLogger.info("Built BM25 index for engine '{}' from {} chunk(s)", this.className, count); + } catch (Exception e) { + classLogger.error("Failed to build BM25 index for engine '{}'; keyword scoring disabled", this.className, e); } } @@ -376,7 +355,7 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT } - // keep the BM25 index in sync with the successful add + // keep the in-memory BM25 index in sync with the successful add if (this.useHybridSearch && this.bm25Index != null && response != null && !response.trim().isEmpty()) { for (int i = 0; i < ids.size(); i++) { Map metadata = metadatas.get(i); @@ -385,8 +364,6 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT this.bm25Index.addRecord(ids.get(i), source == null ? null : source.toString(), content == null ? "" : content.toString(), metadata); } - this.bm25Index.save(); - pushBm25IndexIfClustered(); } return fileStatusList; @@ -411,7 +388,6 @@ public void removeDocument(List fileNames, Map parameter } List filesToRemoveFromCloud = new ArrayList(); - boolean bm25Changed = false; // need to get the source names and then delete it based on the names for (int fileIndex = 0; fileIndex < sourceNames.size(); fileIndex++) { @@ -446,9 +422,9 @@ public void removeDocument(List fileNames, Map parameter this.className); } - // prune the same source from the BM25 index - if (this.useHybridSearch && this.bm25Index != null && this.bm25Index.removeBySource(source) > 0) { - bm25Changed = true; + // prune the same source from the in-memory BM25 index + if (this.useHybridSearch && this.bm25Index != null) { + this.bm25Index.removeBySource(source); } String documentName = Paths.get(fileName).getFileName().toString(); @@ -465,11 +441,6 @@ public void removeDocument(List fileNames, Map parameter } - if (bm25Changed) { - this.bm25Index.save(); - pushBm25IndexIfClustered(); - } - if (ClusterUtil.IS_CLUSTER) { Thread deleteFilesFromCloudThread = new Thread(new DeleteFilesFromEngineRunner(engineId, this.getCatalogType(), filesToRemoveFromCloud.stream().toArray(String[]::new))); From 4b27832566615c6c0259e7b8e40392f83dfb35b4 Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Fri, 26 Jun 2026 11:26:17 -0600 Subject: [PATCH 6/7] fix: resolve bm25 stats race condition and filter matching edge cases in chroma --- .../engine/impl/vector/ChromaBm25Index.java | 53 ++++++++++++------- .../vector/ChromaVectorDatabaseEngine.java | 22 ++++++-- ...omaVectorQueryFilterTranslationHelper.java | 4 ++ 3 files changed, 56 insertions(+), 23 deletions(-) diff --git a/src/prerna/engine/impl/vector/ChromaBm25Index.java b/src/prerna/engine/impl/vector/ChromaBm25Index.java index 3465daa2c3e..3675bee568a 100644 --- a/src/prerna/engine/impl/vector/ChromaBm25Index.java +++ b/src/prerna/engine/impl/vector/ChromaBm25Index.java @@ -29,6 +29,7 @@ import java.util.ArrayList; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.concurrent.locks.ReentrantReadWriteLock; @@ -123,9 +124,9 @@ public List> search(String query, int topK) { return results; } + ensureStats(); lock.readLock().lock(); try { - refreshStatsIfNeeded(); if (records.isEmpty() || avgDocLength <= 0.0) { return results; } @@ -162,21 +163,41 @@ private double scoreRecord(Record record, List queryTerms, int n) { return score; } - /** Recompute document-frequency and average length. Caller must hold a lock. */ - private void refreshStatsIfNeeded() { - if (!statsDirty) { + /** + * Recompute corpus statistics (document frequency, average length) when the index has changed. + * Runs the rebuild under the write lock — never under a read lock — so concurrent searches can't + * race on the shared stats. Double-checked so only one rebuild happens per change. + */ + private void ensureStats() { + lock.readLock().lock(); + boolean dirty; + try { + dirty = statsDirty; + } finally { + lock.readLock().unlock(); + } + if (!dirty) { return; } - docFreq = new LinkedHashMap<>(); - long totalLength = 0; - for (Record record : records.values()) { - totalLength += record.length; - for (String term : record.termFreqs.keySet()) { - docFreq.merge(term, 1, Integer::sum); + lock.writeLock().lock(); + try { + if (!statsDirty) { + return; + } + Map freshDocFreq = new LinkedHashMap<>(); + long totalLength = 0; + for (Record record : records.values()) { + totalLength += record.length; + for (String term : record.termFreqs.keySet()) { + freshDocFreq.merge(term, 1, Integer::sum); + } } + this.docFreq = freshDocFreq; + this.avgDocLength = records.isEmpty() ? 0.0 : (double) totalLength / records.size(); + this.statsDirty = false; + } finally { + lock.writeLock().unlock(); } - avgDocLength = records.isEmpty() ? 0.0 : (double) totalLength / records.size(); - statsDirty = false; } /** Lower-case and split into alphanumeric tokens (no external tokenizer/stemmer). */ @@ -194,12 +215,6 @@ public static List tokenize(String text) { } private static List uniqueTokens(String text) { - List unique = new ArrayList<>(); - for (String token : tokenize(text)) { - if (!unique.contains(token)) { - unique.add(token); - } - } - return unique; + return new ArrayList<>(new LinkedHashSet<>(tokenize(text))); } } diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index ee81272cd2c..4a92310e2e2 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -102,6 +102,12 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private static final double DEFAULT_HYBRID_VECTOR_WEIGHT = 0.5; private static final double DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD = 0.0; private static final int RRF_K = 60; + // over-fetch the hybrid candidate pool: pull more than the caller's limit from each signal + // (vector + keyword) so RRF has a meaningful pool to fuse before truncating back to limit + private static final int HYBRID_CANDIDATE_MULTIPLIER = 10; + private static final int HYBRID_MIN_CANDIDATES = 100; + // the in-memory BM25 index loads the whole collection on open; warn past this size (see PR notes on scale) + private static final int BM25_LARGE_CORPUS_WARN = 100_000; // transient key holding each candidate's vector distance during ranking; stripped before return private static final String DISTANCE_KEY = "_distance"; @@ -109,7 +115,7 @@ public class ChromaVectorDatabaseEngine extends AbstractVectorDatabaseEngine { private double hybridVectorWeight = DEFAULT_HYBRID_VECTOR_WEIGHT; private double hybridKeywordGateThreshold = DEFAULT_HYBRID_KEYWORD_GATE_THRESHOLD; // in-memory keyword index (rebuilt from Chroma on open); used only when hybrid search is enabled - private ChromaBm25Index bm25Index; + private volatile ChromaBm25Index bm25Index; @Override public void open(Properties smssProp) throws Exception { @@ -198,6 +204,12 @@ private void initBm25Index() { content == null ? "" : content.toString(), metadata); } classLogger.info("Built BM25 index for engine '{}' from {} chunk(s)", this.className, count); + if (count >= BM25_LARGE_CORPUS_WARN) { + classLogger.warn( + "BM25 index for engine '{}' loaded {} chunks into memory on open; at this scale consider " + + "paginated loading or native sparse search (Chroma Cloud /search)", + this.className, count); + } } catch (Exception e) { classLogger.error("Failed to build BM25 index for engine '{}'; keyword scoring disabled", this.className, e); } @@ -526,7 +538,7 @@ private void collectFilters(Object value, List target) { */ private List> executeHybridRrfSearch(String searchStatement, List vector, int limit, Map where) { - int candidateLimit = Math.max(limit * 10, 100); + int candidateLimit = Math.max(limit * HYBRID_CANDIDATE_MULTIPLIER, HYBRID_MIN_CANDIDATES); List> vectorHits = queryVectors(vector, candidateLimit, where); List> keywordHits = (this.bm25Index != null) @@ -546,7 +558,9 @@ private List> executeHybridRrfSearch(String searchStatement, distanceById.put(id, distanceOf(hit)); } for (Map hit : keywordHits) { - // the vector side is filtered by Chroma; apply the same filter to the keyword side + // the vector side is filtered by Chroma; apply the same filter to the keyword side. + // NOTE: matches() compares as strings, so it can diverge from Chroma's type coercion for + // non-string metadata (e.g. numeric); our metadata is string-typed so this is consistent today. if (where != null && !ChromaVectorQueryFilterTranslationHelper.matches(hit, where)) { continue; } @@ -747,4 +761,4 @@ private static class ChromaGetResponse { private List> metadatas; } -} \ No newline at end of file +} diff --git a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java index c4aed225fef..85ab0d8a7ad 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java +++ b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java @@ -285,6 +285,10 @@ private static List asList(Object value) { } private static boolean valuesEqual(Object a, Object b) { + // a missing field (null) must not match the literal string "null" + if (a == null || b == null) { + return a == b; + } return String.valueOf(a).equals(String.valueOf(b)); } From a8cce78b17df1948a439ece5b7f6b82493417406 Mon Sep 17 00:00:00 2001 From: Matt Freshwaters Date: Fri, 26 Jun 2026 12:14:10 -0600 Subject: [PATCH 7/7] refactor: recompute bm25 stats per batch under the write lock in chroma --- .../engine/impl/vector/ChromaBm25Index.java | 68 +++++++------------ .../vector/ChromaVectorDatabaseEngine.java | 9 ++- 2 files changed, 32 insertions(+), 45 deletions(-) diff --git a/src/prerna/engine/impl/vector/ChromaBm25Index.java b/src/prerna/engine/impl/vector/ChromaBm25Index.java index 3675bee568a..eded0be2bc2 100644 --- a/src/prerna/engine/impl/vector/ChromaBm25Index.java +++ b/src/prerna/engine/impl/vector/ChromaBm25Index.java @@ -60,7 +60,6 @@ private static class Record { private final Map records = new LinkedHashMap<>(); private Map docFreq = new LinkedHashMap<>(); private double avgDocLength = 0.0; - private boolean statsDirty = true; public boolean isEmpty() { lock.readLock().lock(); @@ -86,7 +85,6 @@ public void addRecord(String id, String source, String content, Map source.equals(r.source)); - int removed = before - records.size(); - if (removed > 0) { - statsDirty = true; + return before - records.size(); + } finally { + lock.writeLock().unlock(); + } + } + + /** + * Recompute corpus statistics (document frequency, average length) under the write lock. The + * engine calls this once after a batch of {@link #addRecord}/{@link #removeBySource} mutations — + * not per record, since a full-corpus load would otherwise be O(n^2). {@link #search} only reads + * these stats, so it never mutates shared state and concurrent searches cannot race. + */ + public void refreshStats() { + lock.writeLock().lock(); + try { + Map freshDocFreq = new LinkedHashMap<>(); + long totalLength = 0; + for (Record record : records.values()) { + totalLength += record.length; + for (String term : record.termFreqs.keySet()) { + freshDocFreq.merge(term, 1, Integer::sum); + } } - return removed; + this.docFreq = freshDocFreq; + this.avgDocLength = records.isEmpty() ? 0.0 : (double) totalLength / records.size(); } finally { lock.writeLock().unlock(); } @@ -124,7 +142,6 @@ public List> search(String query, int topK) { return results; } - ensureStats(); lock.readLock().lock(); try { if (records.isEmpty() || avgDocLength <= 0.0) { @@ -163,43 +180,6 @@ private double scoreRecord(Record record, List queryTerms, int n) { return score; } - /** - * Recompute corpus statistics (document frequency, average length) when the index has changed. - * Runs the rebuild under the write lock — never under a read lock — so concurrent searches can't - * race on the shared stats. Double-checked so only one rebuild happens per change. - */ - private void ensureStats() { - lock.readLock().lock(); - boolean dirty; - try { - dirty = statsDirty; - } finally { - lock.readLock().unlock(); - } - if (!dirty) { - return; - } - lock.writeLock().lock(); - try { - if (!statsDirty) { - return; - } - Map freshDocFreq = new LinkedHashMap<>(); - long totalLength = 0; - for (Record record : records.values()) { - totalLength += record.length; - for (String term : record.termFreqs.keySet()) { - freshDocFreq.merge(term, 1, Integer::sum); - } - } - this.docFreq = freshDocFreq; - this.avgDocLength = records.isEmpty() ? 0.0 : (double) totalLength / records.size(); - this.statsDirty = false; - } finally { - lock.writeLock().unlock(); - } - } - /** Lower-case and split into alphanumeric tokens (no external tokenizer/stemmer). */ public static List tokenize(String text) { List tokens = new ArrayList<>(); diff --git a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java index 4a92310e2e2..a17e5cb0178 100644 --- a/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java +++ b/src/prerna/engine/impl/vector/ChromaVectorDatabaseEngine.java @@ -187,6 +187,7 @@ private void initBm25Index() { include.add("metadatas"); getBody.put("include", include); + // TODO: paginate (Chroma /get supports offset/limit) for very large collections String response = HttpHelperUtility.postRequestStringBody( collection(this.url, this.tenant, this.dbName, this.collectionID, API_GET), buildHeaders(), gson.toJson(getBody), ContentType.APPLICATION_JSON, null, null, null); @@ -203,6 +204,7 @@ private void initBm25Index() { this.bm25Index.addRecord(parsed.ids.get(i), source == null ? null : source.toString(), content == null ? "" : content.toString(), metadata); } + this.bm25Index.refreshStats(); classLogger.info("Built BM25 index for engine '{}' from {} chunk(s)", this.className, count); if (count >= BM25_LARGE_CORPUS_WARN) { classLogger.warn( @@ -376,6 +378,7 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT this.bm25Index.addRecord(ids.get(i), source == null ? null : source.toString(), content == null ? "" : content.toString(), metadata); } + this.bm25Index.refreshStats(); } return fileStatusList; @@ -453,6 +456,10 @@ public void removeDocument(List fileNames, Map parameter } + if (this.useHybridSearch && this.bm25Index != null) { + this.bm25Index.refreshStats(); + } + if (ClusterUtil.IS_CLUSTER) { Thread deleteFilesFromCloudThread = new Thread(new DeleteFilesFromEngineRunner(engineId, this.getCatalogType(), filesToRemoveFromCloud.stream().toArray(String[]::new))); @@ -761,4 +768,4 @@ private static class ChromaGetResponse { private List> metadatas; } -} +}