diff --git a/src/prerna/engine/impl/vector/ChromaBm25Index.java b/src/prerna/engine/impl/vector/ChromaBm25Index.java new file mode 100644 index 0000000000..eded0be2bc --- /dev/null +++ b/src/prerna/engine/impl/vector/ChromaBm25Index.java @@ -0,0 +1,200 @@ +/******************************************************************************* + * 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.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +/** + * 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 double BM25_K1 = 1.2; + private static final double BM25_B = 0.75; + + /** 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 source; + private Map metadata; + private Map termFreqs; + private int length; + } + + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private final Map records = new LinkedHashMap<>(); + private Map docFreq = new LinkedHashMap<>(); + private double avgDocLength = 0.0; + + 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.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); + } 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)); + 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); + } + } + this.docFreq = freshDocFreq; + this.avgDocLength = records.isEmpty() ? 0.0 : (double) totalLength / records.size(); + } 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 { + if (records.isEmpty() || avgDocLength <= 0.0) { + return results; + } + int n = records.size(); + for (Map.Entry entry : records.entrySet()) { + double score = scoreRecord(entry.getValue(), queryTerms, n); + if (score <= 0.0) { + continue; + } + Map row = new LinkedHashMap<>(entry.getValue().metadata); + row.put(ID_KEY, entry.getKey()); + 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; + } + + /** 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) { + 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 d640e3fb4c..a17e5cb017 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; @@ -46,7 +47,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; @@ -54,6 +54,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 +67,56 @@ 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 final String API_GET = "/get"; + 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: enable hybrid (vector + keyword) search. */ + public static final String USE_HYBRID_SEARCH = "USE_HYBRID_SEARCH"; + /** 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: 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; + 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"; + + private boolean useHybridSearch = false; + 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 volatile ChromaBm25Index bm25Index; + @Override public void open(Properties smssProp) throws Exception { super.open(smssProp); @@ -87,32 +127,135 @@ 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")); + + // 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); + } + } + + 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); + + if (this.useHybridSearch) { + initBm25Index(); + } } /** - * - * @param collectionName + * 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() { + this.bm25Index = new ChromaBm25Index(); + try { + Gson gson = new Gson(); + Map getBody = new HashMap<>(); + List include = new ArrayList<>(); + 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); + + 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.refreshStats(); + 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); + } + } + + /** 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; + } + 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(); + } + + /** 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(); + } + + /** 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; + } + Map headers = new HashMap<>(); + headers.put(API_TOKEN_KEY, this.apiKey); + headers.put(HttpHeaders.CONTENT_TYPE, ContentType.APPLICATION_JSON.getMimeType()); + return headers; + } + + /** 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 GsonBuilder().setPrettyPrinting().create(); - 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; - } - + Gson gson = new Gson(); + String collectionsUrl = collections(this.url, this.tenant, this.dbName); + Map headersMap = buildHeaders(); + 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 +273,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"); } @@ -196,16 +339,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(this.url + this.collectionID + API_ADD, - headersMap, body, ContentType.APPLICATION_JSON, null, null, null); + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_ADD), + 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()) { @@ -233,6 +369,18 @@ public List addEmbeddings(VectorDatabaseCSVTable vectorCsvT } + // 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); + 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.refreshStats(); + } + return fileStatusList; } @@ -260,10 +408,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 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 chunks matching the Source metadata (v2 delete API) Map fileNamesForDelete = new HashMap<>(); Map sourceProperty = new HashMap<>(); @@ -273,21 +418,30 @@ 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); - 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()); + String response = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_DELETE), + buildHeaders(), body, ContentType.APPLICATION_JSON, null, null, null); + + // 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 { - headersMap = null; + classLogger.warn("No records matched source '{}' in Chroma collection '{}' during delete", source, + this.className); } - String response = HttpHelperUtility.postRequestStringBody(this.url + this.collectionID + API_DELETE, - headersMap, body, ContentType.APPLICATION_JSON, null, null, null); + // prune the same source from the in-memory BM25 index + if (this.useHybridSearch && this.bm25Index != null) { + this.bm25Index.removeBySource(source); + } - //TODO: let us add validation by looking at the response - 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); @@ -302,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))); @@ -310,7 +468,8 @@ public void removeDocument(List fileNames, Map parameter } @Override - public List> nearestNeighborCall(Insight insight, String searchStatement, Number limit, Map parameters) { + 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,40 +479,229 @@ public List> nearestNeighborCall(Insight insight, String sea if (limit == null) { limit = 3; } - - Gson gson = new Gson(); List vector = getEmbeddingsDouble(searchStatement, insight); + Map where = buildWhereClause(parameters); + + if (this.useHybridSearch) { + return executeHybridRrfSearch(searchStatement, vector, limit.intValue(), where); + } + + // 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); + candidate.remove(ChromaBm25Index.ID_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; + } + + /** 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; + } + return 1.0 / (1.0 + distance); + } + + /** Build a Chroma {@code where} clause from the {@code filters}/{@code metaFilters} params (AND-combined), or {@code null}. */ + private Map buildWhereClause(Map parameters) { + if (parameters == null) { + return null; + } + List allFilters = new ArrayList<>(); + collectFilters(parameters.get(AbstractVectorDatabaseEngine.FILTERS_KEY), allFilters); + collectFilters(parameters.get(AbstractVectorDatabaseEngine.METADATA_FILTERS_KEY), allFilters); + return ChromaVectorQueryFilterTranslationHelper.toWhere(allFilters); + } + + /** 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; + } + for (Object element : (List) value) { + if (element instanceof IQueryFilter) { + target.add((IQueryFilter) element); + } + } + } + + /** + * 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 * HYBRID_CANDIDATE_MULTIPLIER, HYBRID_MIN_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. + // 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; + } + 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<>(); + } + + // 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 '{}': 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 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++) { + 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; + } + + /** 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; + } + + private static double maxValue(Map values) { + double max = 0.0; + for (double value : values.values()) { + if (value > max) { + max = value; + } + } + return max; + } + + /** + * 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(); 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); - query.put("n_results", limit); + // nest the embedding inside a list as the API expects + queryEmbeddings.add(vector); + query.put("n_results", nResults); query.put("query_embeddings", queryEmbeddings); - String body = gson.toJson(query); + if (where != null) { + query.put("where", where); + } - 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 responseBody = HttpHelperUtility.postRequestStringBody( + collection(this.url, this.tenant, this.dbName, this.collectionID, API_QUERY), + buildHeaders(), gson.toJson(query), ContentType.APPLICATION_JSON, null, null, 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<>(); } - - String nearestNeigborResponse = HttpHelperUtility.postRequestStringBody(this.url + 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); + 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; } - + + /** 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; + } + @Override public List> listDocuments(Map parameters) { //TODO: needs to grab 'Source' from the database @@ -402,4 +750,22 @@ public VectorDatabaseTypeEnum getVectorDatabaseType() { return VectorDatabaseTypeEnum.CHROMA; } -} \ No newline at end of file + /** 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; + } + + /** Typed Chroma delete response: {@code {"deleted": N}}. */ + 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; + } + +} diff --git a/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java new file mode 100644 index 0000000000..85ab0d8a7a --- /dev/null +++ b/src/prerna/engine/impl/vector/ChromaVectorQueryFilterTranslationHelper.java @@ -0,0 +1,303 @@ +/******************************************************************************* + * 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 — 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 { + + private static final String AND = "$and"; + private static final String OR = "$or"; + + private ChromaVectorQueryFilterTranslationHelper() { + } + + /** 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; + } + 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 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; + } + 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 map 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; + } + + /** Column name from a simple filter's left comparison. */ + 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(); + } + + /** Flatten a right-comparison value to a list (handles a single value or a collection). */ + 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; + } + + /** + * 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) { + // 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)); + } + + private static boolean containsValue(List values, Object actual) { + for (Object value : values) { + if (valuesEqual(actual, value)) { + return true; + } + } + return false; + } +}