From 2981db84847da8f4ace735ca7b456692fc9b834b Mon Sep 17 00:00:00 2001 From: Kaveh-Vakili Date: Wed, 12 Aug 2026 10:05:01 -0400 Subject: [PATCH] feat: add cache-aware "Token + Cache" usage limit type Adds a new per-model, cache-token-aware usage limit alongside the existing Token/Compute time restrictions. Admins can configure a cache read/write weight (0-1000%) per model via Model Settings or the SMSS file, and members can be restricted with a new "Token + Cache" limit that counts cache read/creation tokens toward their token budget at that weighted percentage, on top of normal prompt+completion tokens. - Constants: add CACHE_READ_WEIGHT/CACHE_WRITE_WEIGHT metadata keys and MODEL_TOKEN_CACHE_RESTRICTION_VALUE restriction value - SecurityOwlCreator: declare CACHEREADWEIGHT/CACHEWRITEWEIGHT columns on MODELMETADATA - SecurityModelMetadataUtils: validate (0-1000%), persist, and seed the two new fields from smss properties - ModelUsageRestrictionUtility: new token_cache branch (engine- and user-level) that resolves a model's configured weights and folds weighted cache tokens into the usage check and running total - ModelInferenceLogsUtils: sum CACHE_READ_TOKENS/CACHE_CREATION_TOKENS alongside MESSAGE_TOKENS for token-family restriction modes - AbstractModelEngineResponse: carry resolved weights between the limit check and the per-call update Also fixes a pre-existing bug found while testing: getEngineUsagePermissionMap's SMSS_USER/ENGINEPERMISSION left outer join used bare table names and was silently resolving to zero rows for every restriction mode, not just this one. Switched to explicit qualified join columns, matching the existing pattern in SecurityProjectUtils. Note: MODELMETADATA is a pre-existing table; the OWL schema declaration above does not auto-ALTER existing databases, so any environment upgrading onto an existing security DB needs the two columns added manually (ALTER TABLE MODELMETADATA ADD COLUMN CACHEREADWEIGHT DOUBLE / CACHEWRITEWEIGHT DOUBLE) before this will work. --- .../auth/utils/SecurityEngineUtils.java | 6 +- .../utils/SecurityModelMetadataUtils.java | 81 ++++++++++++-- src/prerna/auth/utils/SecurityOwlCreator.java | 4 +- .../model/ModelUsageRestrictionUtility.java | 105 ++++++++++++++++-- .../ModelInferenceLogsUtils.java | 85 ++++++++++---- .../AbstractModelEngineResponse.java | 2 + src/prerna/util/Constants.java | 8 ++ 7 files changed, 252 insertions(+), 39 deletions(-) diff --git a/src/prerna/auth/utils/SecurityEngineUtils.java b/src/prerna/auth/utils/SecurityEngineUtils.java index 16f7e20a06d..06471715d56 100644 --- a/src/prerna/auth/utils/SecurityEngineUtils.java +++ b/src/prerna/auth/utils/SecurityEngineUtils.java @@ -3537,8 +3537,10 @@ public static List> getEngineUsagePermissionMap(User user, S qs.addExplicitFilter( SimpleQueryFilter.makeColToValFilter("ENGINEPERMISSION__USERID", "==", userDetails.getValue0())); qs.addExplicitFilter(SimpleQueryFilter.makeColToValFilter("ENGINEPERMISSION__ENGINEID", "==", engineId)); - // relationship between SMSS_USER and ENGINEPERMISSION tables - qs.addRelation("SMSS_USER", "ENGINEPERMISSION", "left.outer.join"); + // relationship between SMSS_USER and ENGINEPERMISSION tables. Explicit + // join columns (rather than bare table names) since auto-resolution was + // silently returning zero rows for this left outer join. + qs.addRelation("SMSS_USER__ID", "ENGINEPERMISSION__USERID", "left.outer.join"); return QueryExecutionUtility.flushRsToMap(securityDb, qs); } diff --git a/src/prerna/auth/utils/SecurityModelMetadataUtils.java b/src/prerna/auth/utils/SecurityModelMetadataUtils.java index e2016b650a0..50a91dca312 100644 --- a/src/prerna/auth/utils/SecurityModelMetadataUtils.java +++ b/src/prerna/auth/utils/SecurityModelMetadataUtils.java @@ -82,13 +82,15 @@ public final class SecurityModelMetadataUtils extends AbstractSecurityUtils { private static final Set EDITABLE_METADATA_KEYS = Set.of(Constants.MODEL_PROVIDER, Constants.SERVING_PROVIDER, Constants.MODEL_CAPABILITY, Constants.INPUT_MODALITIES, Constants.OUTPUT_MODALITIES, Constants.CONTEXT_WINDOW, Constants.MAX_TOKENS, Constants.BUILTIN_TOOLS, - Constants.REASONING, Constants.REASONING_CONFIG); + Constants.REASONING, Constants.REASONING_CONFIG, Constants.CACHE_READ_WEIGHT, + Constants.CACHE_WRITE_WEIGHT); private static final Set CATALOG_ONLY_KEYS = Set.of(Constants.MODEL_PROVIDER, Constants.SERVING_PROVIDER, Constants.MODEL_CAPABILITY, Constants.INPUT_MODALITIES, Constants.OUTPUT_MODALITIES, Constants.BUILTIN_TOOLS, Constants.MODEL_FAMILY, Constants.ATTACHMENT, Constants.REASONING, Constants.TOOL_CALL, Constants.STRUCTURED_OUTPUT, Constants.TEMPERATURE, Constants.KNOWLEDGE_CUTOFF, Constants.RELEASE_DATE, Constants.SUPPORTED_PARAMETERS, - Constants.REASONING_CONFIG, Constants.BENCHMARKS, Constants.DESCR); + Constants.REASONING_CONFIG, Constants.BENCHMARKS, Constants.DESCR, Constants.CACHE_READ_WEIGHT, + Constants.CACHE_WRITE_WEIGHT); private static final Set REMOVED_METADATA_KEYS = Set.of("LICENSE", "LINKS", "WEIGHTS", "OPEN_WEIGHTS", "LAST_UPDATED", Constants.MAX_INPUT_TOKENS); private static final Pattern IDENTIFIER_PATTERN = Pattern.compile("^[A-Z][A-Z0-9_]*$"); @@ -129,6 +131,8 @@ public static Map normalizeModelDetails(Map mode normalizeJsonArrayProperty(normalized, Constants.BENCHMARKS); normalizePositiveLongProperty(normalized, Constants.CONTEXT_WINDOW); normalizePositiveLongProperty(normalized, Constants.MAX_TOKENS); + normalizeWeightPercentageProperty(normalized, Constants.CACHE_READ_WEIGHT); + normalizeWeightPercentageProperty(normalized, Constants.CACHE_WRITE_WEIGHT); return normalized; } @@ -182,6 +186,8 @@ public static void upsertModelMetadata(String engineId, Properties properties) { copyIfPresent(properties, details, Constants.SUPPORTED_PARAMETERS); copyIfPresent(properties, details, Constants.REASONING_CONFIG); copyIfPresent(properties, details, Constants.BENCHMARKS); + copyIfPresent(properties, details, Constants.CACHE_READ_WEIGHT); + copyIfPresent(properties, details, Constants.CACHE_WRITE_WEIGHT); Map merged = toDetails(getModelMetadata(engineId)); merged.putAll(details); @@ -205,8 +211,8 @@ public static void upsertModelMetadata(String engineId, Map mode IRDBMSEngine securityDb = SystemEngineRegistry.getSecurityDb(); boolean exists = modelMetadataExists(securityDb, metadata.engineId()); String sql = exists - ? "UPDATE MODELMETADATA SET MODELID=?, MODELPROVIDER=?, SERVINGPROVIDER=?, CAPABILITY=?, FAMILY=?, INPUTMODALITIES=?, OUTPUTMODALITIES=?, CONTEXTWINDOW=?, MAXOUTPUTTOKENS=?, BUILTINTOOLS=?, ATTACHMENT=?, REASONING=?, TOOLCALL=?, STRUCTUREDOUTPUT=?, TEMPERATURE=?, KNOWLEDGECUTOFF=?, RELEASEDATE=?, SUPPORTEDPARAMETERS=?, REASONINGCONFIG=?, BENCHMARKS=? WHERE ENGINEID=?" - : "INSERT INTO MODELMETADATA (MODELID, MODELPROVIDER, SERVINGPROVIDER, CAPABILITY, FAMILY, INPUTMODALITIES, OUTPUTMODALITIES, CONTEXTWINDOW, MAXOUTPUTTOKENS, BUILTINTOOLS, ATTACHMENT, REASONING, TOOLCALL, STRUCTUREDOUTPUT, TEMPERATURE, KNOWLEDGECUTOFF, RELEASEDATE, SUPPORTEDPARAMETERS, REASONINGCONFIG, BENCHMARKS, ENGINEID) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"; + ? "UPDATE MODELMETADATA SET MODELID=?, MODELPROVIDER=?, SERVINGPROVIDER=?, CAPABILITY=?, FAMILY=?, INPUTMODALITIES=?, OUTPUTMODALITIES=?, CONTEXTWINDOW=?, MAXOUTPUTTOKENS=?, BUILTINTOOLS=?, ATTACHMENT=?, REASONING=?, TOOLCALL=?, STRUCTUREDOUTPUT=?, TEMPERATURE=?, KNOWLEDGECUTOFF=?, RELEASEDATE=?, SUPPORTEDPARAMETERS=?, REASONINGCONFIG=?, BENCHMARKS=?, CACHEREADWEIGHT=?, CACHEWRITEWEIGHT=? WHERE ENGINEID=?" + : "INSERT INTO MODELMETADATA (MODELID, MODELPROVIDER, SERVINGPROVIDER, CAPABILITY, FAMILY, INPUTMODALITIES, OUTPUTMODALITIES, CONTEXTWINDOW, MAXOUTPUTTOKENS, BUILTINTOOLS, ATTACHMENT, REASONING, TOOLCALL, STRUCTUREDOUTPUT, TEMPERATURE, KNOWLEDGECUTOFF, RELEASEDATE, SUPPORTEDPARAMETERS, REASONINGCONFIG, BENCHMARKS, CACHEREADWEIGHT, CACHEWRITEWEIGHT, ENGINEID) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"; PreparedStatement ps = null; try { @@ -232,6 +238,8 @@ public static void upsertModelMetadata(String engineId, Map mode setNullableString(ps, index++, metadata.supportedParametersJson()); setNullableString(ps, index++, metadata.reasoningConfigJson()); setNullableString(ps, index++, metadata.benchmarksJson()); + setNullableDouble(ps, index++, metadata.cacheReadWeight()); + setNullableDouble(ps, index++, metadata.cacheWriteWeight()); ps.setString(index, metadata.engineId()); ps.executeUpdate(); ConnectionUtils.commitConnection(ps.getConnection()); @@ -449,6 +457,8 @@ private static Map toDetails(Map existing) { details.put(Constants.SUPPORTED_PARAMETERS, existing.get("supportedParameters")); details.put(Constants.REASONING_CONFIG, existing.get("reasoningConfig")); details.put(Constants.BENCHMARKS, existing.get("benchmarks")); + details.put(Constants.CACHE_READ_WEIGHT, existing.get("cacheReadWeight")); + details.put(Constants.CACHE_WRITE_WEIGHT, existing.get("cacheWriteWeight")); return details; } @@ -457,7 +467,7 @@ private static Map toDetails(Map existing) { */ public static Map getModelMetadata(String engineId) { IRDBMSEngine securityDb = SystemEngineRegistry.getSecurityDb(); - String sql = "SELECT ENGINEID, MODELID, MODELPROVIDER, SERVINGPROVIDER, CAPABILITY, FAMILY, INPUTMODALITIES, OUTPUTMODALITIES, CONTEXTWINDOW, MAXOUTPUTTOKENS, BUILTINTOOLS, ATTACHMENT, REASONING, TOOLCALL, STRUCTUREDOUTPUT, TEMPERATURE, KNOWLEDGECUTOFF, RELEASEDATE, SUPPORTEDPARAMETERS, REASONINGCONFIG, BENCHMARKS FROM MODELMETADATA WHERE ENGINEID=?"; + String sql = "SELECT ENGINEID, MODELID, MODELPROVIDER, SERVINGPROVIDER, CAPABILITY, FAMILY, INPUTMODALITIES, OUTPUTMODALITIES, CONTEXTWINDOW, MAXOUTPUTTOKENS, BUILTINTOOLS, ATTACHMENT, REASONING, TOOLCALL, STRUCTUREDOUTPUT, TEMPERATURE, KNOWLEDGECUTOFF, RELEASEDATE, SUPPORTEDPARAMETERS, REASONINGCONFIG, BENCHMARKS, CACHEREADWEIGHT, CACHEWRITEWEIGHT FROM MODELMETADATA WHERE ENGINEID=?"; PreparedStatement ps = null; ResultSet rs = null; try { @@ -499,7 +509,7 @@ public static Map> getModelMetadata(Collection batch = normalizedEngineIds.subList(start, end); String placeholders = String.join(",", Collections.nCopies(batch.size(), "?")); - String sql = "SELECT ENGINEID, MODELID, MODELPROVIDER, SERVINGPROVIDER, CAPABILITY, FAMILY, INPUTMODALITIES, OUTPUTMODALITIES, CONTEXTWINDOW, MAXOUTPUTTOKENS, BUILTINTOOLS, ATTACHMENT, REASONING, TOOLCALL, STRUCTUREDOUTPUT, TEMPERATURE, KNOWLEDGECUTOFF, RELEASEDATE, SUPPORTEDPARAMETERS, REASONINGCONFIG, BENCHMARKS FROM MODELMETADATA WHERE ENGINEID IN (" + String sql = "SELECT ENGINEID, MODELID, MODELPROVIDER, SERVINGPROVIDER, CAPABILITY, FAMILY, INPUTMODALITIES, OUTPUTMODALITIES, CONTEXTWINDOW, MAXOUTPUTTOKENS, BUILTINTOOLS, ATTACHMENT, REASONING, TOOLCALL, STRUCTUREDOUTPUT, TEMPERATURE, KNOWLEDGECUTOFF, RELEASEDATE, SUPPORTEDPARAMETERS, REASONINGCONFIG, BENCHMARKS, CACHEREADWEIGHT, CACHEWRITEWEIGHT FROM MODELMETADATA WHERE ENGINEID IN (" + placeholders + ")"; PreparedStatement ps = null; @@ -554,6 +564,8 @@ public static Map toCapabilities(Map modelMetada capabilities.put("supportedParameters", emptyListIfNull(modelMetadata.get("supportedParameters"))); capabilities.put("reasoningConfig", emptyMapIfNull(modelMetadata.get("reasoningConfig"))); capabilities.put("benchmarks", emptyListIfNull(modelMetadata.get("benchmarks"))); + capabilities.put("cacheReadWeight", emptyStringIfNull(modelMetadata.get("cacheReadWeight"))); + capabilities.put("cacheWriteWeight", emptyStringIfNull(modelMetadata.get("cacheWriteWeight"))); return capabilities; } @@ -584,7 +596,8 @@ private static boolean containsMetadata(Map details) { || details.containsKey(Constants.TOOL_CALL) || details.containsKey(Constants.STRUCTURED_OUTPUT) || details.containsKey(Constants.TEMPERATURE) || details.containsKey(Constants.KNOWLEDGE_CUTOFF) || details.containsKey(Constants.RELEASE_DATE) || details.containsKey(Constants.SUPPORTED_PARAMETERS) - || details.containsKey(Constants.REASONING_CONFIG) || details.containsKey(Constants.BENCHMARKS); + || details.containsKey(Constants.REASONING_CONFIG) || details.containsKey(Constants.BENCHMARKS) + || details.containsKey(Constants.CACHE_READ_WEIGHT) || details.containsKey(Constants.CACHE_WRITE_WEIGHT); } /** @@ -618,7 +631,9 @@ private static ModelMetadata toMetadata(String engineId, Map det nullableString(details.get(Constants.KNOWLEDGE_CUTOFF)), nullableString(details.get(Constants.RELEASE_DATE)), nullableString(details.get(Constants.SUPPORTED_PARAMETERS)), - nullableString(details.get(Constants.REASONING_CONFIG)), nullableString(details.get(Constants.BENCHMARKS))); + nullableString(details.get(Constants.REASONING_CONFIG)), nullableString(details.get(Constants.BENCHMARKS)), + toNullableDouble(details.get(Constants.CACHE_READ_WEIGHT)), + toNullableDouble(details.get(Constants.CACHE_WRITE_WEIGHT))); } private static void normalizeStringProperty(Map details, String key, boolean identifier) { @@ -832,6 +847,8 @@ private static Map readModelMetadata(ResultSet rs) throws SQLExc metadata.put("supportedParameters", parseStoredList(rs.getString("SUPPORTEDPARAMETERS"))); metadata.put("reasoningConfig", parseStoredJsonObject(rs.getString("REASONINGCONFIG"))); metadata.put("benchmarks", parseStoredJsonArray(rs.getString("BENCHMARKS"))); + metadata.put("cacheReadWeight", getNullableDouble(rs, "CACHEREADWEIGHT")); + metadata.put("cacheWriteWeight", getNullableDouble(rs, "CACHEWRITEWEIGHT")); return metadata; } @@ -848,6 +865,24 @@ private static void normalizePositiveLongProperty(Map details, S details.put(key, value == null ? "" : value); } + /** + * Validate a cache token weight: a percentage of a normal token this cache + * token counts as toward a member's "token" usage restriction. Zero is valid + * (e.g. cache reads counting for nothing), unlike + * {@link #normalizePositiveLongProperty}; the 1000% ceiling is generous + * headroom against fat-fingered entry rather than a real provider limit. + */ + private static void normalizeWeightPercentageProperty(Map details, String key) { + if (!details.containsKey(key)) { + return; + } + Double value = toNullableDouble(details.get(key)); + if (value != null && (value < 0 || value > 1000)) { + throw new IllegalArgumentException(key + " must be between 0 and 1000"); + } + details.put(key, value == null ? "" : value); + } + private static Long toNullableLong(Object value) { if (value == null || value.toString().trim().isEmpty()) { return null; @@ -866,6 +901,20 @@ private static Long toNullableLong(Object value) { } } + private static Double toNullableDouble(Object value) { + if (value == null || value.toString().trim().isEmpty()) { + return null; + } + if (value instanceof Number number) { + return number.doubleValue(); + } + try { + return Double.valueOf(value.toString().trim()); + } catch (NumberFormatException e) { + throw new IllegalArgumentException("Expected a number but received " + value, e); + } + } + private static Boolean toNullableBoolean(Object value) { if (value == null || value.toString().trim().isEmpty()) { return null; @@ -942,11 +991,24 @@ private static void setNullableBoolean(PreparedStatement ps, int index, Boolean } } + private static void setNullableDouble(PreparedStatement ps, int index, Double value) throws SQLException { + if (value == null) { + ps.setNull(index, Types.DOUBLE); + } else { + ps.setDouble(index, value); + } + } + private static Long getNullableLong(ResultSet rs, String column) throws SQLException { long value = rs.getLong(column); return rs.wasNull() ? null : value; } + private static Double getNullableDouble(ResultSet rs, String column) throws SQLException { + double value = rs.getDouble(column); + return rs.wasNull() ? null : value; + } + private static Boolean getNullableBoolean(ResultSet rs, String column) throws SQLException { boolean value = rs.getBoolean(column); return rs.wasNull() ? null : value; @@ -956,6 +1018,7 @@ private record ModelMetadata(String engineId, String modelId, String modelProvid String capability, String family, String inputModalitiesJson, String outputModalitiesJson, Long contextWindow, Long maxOutputTokens, String builtinToolsJson, Boolean attachment, Boolean reasoning, Boolean toolCall, Boolean structuredOutput, Boolean temperature, String knowledgeCutoff, String releaseDate, - String supportedParametersJson, String reasoningConfigJson, String benchmarksJson) { + String supportedParametersJson, String reasoningConfigJson, String benchmarksJson, Double cacheReadWeight, + Double cacheWriteWeight) { } } diff --git a/src/prerna/auth/utils/SecurityOwlCreator.java b/src/prerna/auth/utils/SecurityOwlCreator.java index 5c33583f1f4..45070e4cf95 100644 --- a/src/prerna/auth/utils/SecurityOwlCreator.java +++ b/src/prerna/auth/utils/SecurityOwlCreator.java @@ -107,7 +107,9 @@ public void createColumnsAndTypes(AbstractSqlQueryUtil queryUtil) { Pair.with("RELEASEDATE", VARCHAR_255), Pair.with("SUPPORTEDPARAMETERS", CLOB_DATATYPE_NAME), Pair.with("REASONINGCONFIG", CLOB_DATATYPE_NAME), - Pair.with("BENCHMARKS", CLOB_DATATYPE_NAME))); + Pair.with("BENCHMARKS", CLOB_DATATYPE_NAME), + Pair.with("CACHEREADWEIGHT", DOUBLE_DATATYPE_NAME), + Pair.with("CACHEWRITEWEIGHT", DOUBLE_DATATYPE_NAME))); addTable("ENGINEPERMISSION", Arrays.asList( Pair.with("ENGINEID", VARCHAR_255), diff --git a/src/prerna/engine/impl/model/ModelUsageRestrictionUtility.java b/src/prerna/engine/impl/model/ModelUsageRestrictionUtility.java index 8f777160998..5ae83e40ba7 100644 --- a/src/prerna/engine/impl/model/ModelUsageRestrictionUtility.java +++ b/src/prerna/engine/impl/model/ModelUsageRestrictionUtility.java @@ -41,6 +41,7 @@ import prerna.auth.User; import prerna.auth.utils.SecurityEngineUtils; +import prerna.auth.utils.SecurityModelMetadataUtils; import prerna.engine.impl.model.inferencetracking.ModelInferenceLogsUtils; import prerna.engine.impl.model.responses.AbstractModelEngineResponse; import prerna.util.Constants; @@ -56,8 +57,31 @@ public final class ModelUsageRestrictionUtility { public static final String ENGINE_TOKEN_LIMIT_EXCEEDED_MESSAGE = "Token limit exceeded for engine level: You have used %d tokens, but the limit is %d"; public static final String ENGINE_RESPONSE_TIME_LIMIT_EXCEEDED_MESSAGE = "Response time limit exceeded for engine level : You have reached %.2f seconds, but the limit is %.2f seconds."; + // Applied when a model has never had a cache weight configured, so counting + // cache tokens toward a "token" restriction starts out equivalent to counting + // them as regular tokens rather than silently ignoring them. + private static final double DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT = 100.0; + /** - * + * Look up the admin-configured cache token weights for a model engine, in + * [cacheReadWeightPercent, cacheWriteWeightPercent] order. Falls back to + * {@link #DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT} for either value that has never + * been set. + * + * @param engineId the model engine + * @return the two weight percentages + */ + private static double[] resolveCacheTokenWeightPercents(String engineId) { + Map metadata = SecurityModelMetadataUtils.getModelMetadata(engineId); + Object cacheReadWeight = metadata == null ? null : metadata.get("cacheReadWeight"); + Object cacheWriteWeight = metadata == null ? null : metadata.get("cacheWriteWeight"); + return new double[] { + cacheReadWeight instanceof Number number ? number.doubleValue() : DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT, + cacheWriteWeight instanceof Number number ? number.doubleValue() : DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT }; + } + + /** + * * @param user * @param engineId * @return @@ -102,7 +126,7 @@ public static Map getModelUsageRestriction(User user, String eng if (Constants.MODEL_TOKEN_RESTRICTION_VALUE.equalsIgnoreCase(engineLvlModelUsageRestriction)) { currentUsage = ModelInferenceLogsUtils.getTotalTokensOrTotalResponseTime( Constants.MODEL_TOKEN_RESTRICTION_VALUE, user, engineId, currentDateTime, - engineLvlModelUsageFrequency); + engineLvlModelUsageFrequency, 0.0, 0.0); if (currentUsage.intValue() > engineLvlModelUsageMaxTokens.intValue()) { throw new IllegalArgumentException(String.format(ENGINE_TOKEN_LIMIT_EXCEEDED_MESSAGE, @@ -116,11 +140,35 @@ public static Map getModelUsageRestriction(User user, String eng userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_MAX_VALUE, engineLvlModelUsageMaxTokens.intValue()); + } else if (Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE + .equalsIgnoreCase(engineLvlModelUsageRestriction)) { + double[] cacheWeights = resolveCacheTokenWeightPercents(engineId); + currentUsage = ModelInferenceLogsUtils.getTotalTokensOrTotalResponseTime( + Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE, user, engineId, currentDateTime, + engineLvlModelUsageFrequency, cacheWeights[0], cacheWeights[1]); + + if (currentUsage.intValue() > engineLvlModelUsageMaxTokens.intValue()) { + throw new IllegalArgumentException(String.format(ENGINE_TOKEN_LIMIT_EXCEEDED_MESSAGE, + currentUsage.intValue(), engineLvlModelUsageMaxTokens.intValue())); + } + + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_MODE, + Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CURRENT_VALUE, + currentUsage.intValue()); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_MAX_VALUE, + engineLvlModelUsageMaxTokens.intValue()); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CACHE_READ_WEIGHT, + cacheWeights[0]); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CACHE_WRITE_WEIGHT, + cacheWeights[1]); + } else if (Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE .equalsIgnoreCase(engineLvlModelUsageRestriction)) { currentUsage = ModelInferenceLogsUtils.getTotalTokensOrTotalResponseTime( Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE, user, engineId, currentDateTime, - engineLvlModelUsageFrequency); + engineLvlModelUsageFrequency, DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT, + DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT); if (currentUsage.doubleValue() > engineLvlModelUsageMaxResponseTime.doubleValue()) { throw new IllegalArgumentException(String.format(ENGINE_RESPONSE_TIME_LIMIT_EXCEEDED_MESSAGE, @@ -149,7 +197,7 @@ else if (userLvlModelUsageRestriction != null && !userLvlModelUsageRestriction.i if (Constants.MODEL_TOKEN_RESTRICTION_VALUE.equalsIgnoreCase(userLvlModelUsageRestriction)) { currentUsage = ModelInferenceLogsUtils.getTotalUsageForUser(Constants.MODEL_TOKEN_RESTRICTION_VALUE, - user, engineId, currentDateTime, userLvlModelUsageFrequency); + user, engineId, currentDateTime, userLvlModelUsageFrequency, 0.0, 0.0); if (currentUsage.intValue() > userLvlModelUsageMaxTokens.intValue()) { throw new IllegalArgumentException(String.format(USER_TOKEN_LIMIT_EXCEEDED_MESSAGE, @@ -162,12 +210,36 @@ else if (userLvlModelUsageRestriction != null && !userLvlModelUsageRestriction.i userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_MAX_VALUE, userLvlModelUsageMaxTokens.intValue()); + } else if (Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE + .equalsIgnoreCase(userLvlModelUsageRestriction)) { + + double[] cacheWeights = resolveCacheTokenWeightPercents(engineId); + currentUsage = ModelInferenceLogsUtils.getTotalUsageForUser( + Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE, user, engineId, currentDateTime, + userLvlModelUsageFrequency, cacheWeights[0], cacheWeights[1]); + + if (currentUsage.intValue() > userLvlModelUsageMaxTokens.intValue()) { + throw new IllegalArgumentException(String.format(USER_TOKEN_LIMIT_EXCEEDED_MESSAGE, + currentUsage.intValue(), userLvlModelUsageMaxTokens.intValue())); + } + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_MODE, + Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CURRENT_VALUE, + currentUsage.intValue()); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_MAX_VALUE, + userLvlModelUsageMaxTokens.intValue()); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CACHE_READ_WEIGHT, + cacheWeights[0]); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CACHE_WRITE_WEIGHT, + cacheWeights[1]); + } else if (Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE .equalsIgnoreCase(userLvlModelUsageRestriction)) { currentUsage = ModelInferenceLogsUtils.getTotalUsageForUser( Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE, user, engineId, currentDateTime, - userLvlModelUsageFrequency); + userLvlModelUsageFrequency, DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT, + DEFAULT_CACHE_TOKEN_WEIGHT_PERCENT); if (currentUsage.doubleValue() > userLvlModelUsageMaxResponseTime.doubleValue()) { throw new IllegalArgumentException(String.format(USER_RESPONSE_TIME_LIMIT_EXCEEDED_MESSAGE, @@ -203,13 +275,30 @@ public static void updateRestrictionMapCurrentUsage(Map userRest String restrictionMode = (String) userRestrictionMap .get(AbstractModelEngineResponse.USAGE_RESTRICTION_MODE); - if (Constants.MODEL_TOKEN_RESTRICTION_VALUE.equalsIgnoreCase(restrictionMode)) { + if (Constants.MODEL_TOKEN_RESTRICTION_VALUE.equalsIgnoreCase(restrictionMode) + || Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE.equalsIgnoreCase(restrictionMode)) { + // plain "token" mode never stashes a weight, so 0 (i.e. cache tokens do not + // count) is the correct default there; "token_cache" mode always stashes the + // model's resolved weight (itself defaulting to 100 when unconfigured) + double cacheReadWeightPercent = ((Number) userRestrictionMap + .getOrDefault(AbstractModelEngineResponse.USAGE_RESTRICTION_CACHE_READ_WEIGHT, 0.0)) + .doubleValue(); + double cacheWriteWeightPercent = ((Number) userRestrictionMap + .getOrDefault(AbstractModelEngineResponse.USAGE_RESTRICTION_CACHE_WRITE_WEIGHT, 0.0)) + .doubleValue(); + int cacheReadTokens = modelResponse.getNumberOfCacheReadTokens() == null ? 0 + : modelResponse.getNumberOfCacheReadTokens(); + int cacheCreationTokens = modelResponse.getNumberOfCacheCreationTokens() == null ? 0 + : modelResponse.getNumberOfCacheCreationTokens(); + double weightedCacheTokens = cacheReadTokens * (cacheReadWeightPercent / 100.0) + + cacheCreationTokens * (cacheWriteWeightPercent / 100.0); + userRestrictionMap.put(AbstractModelEngineResponse.USAGE_RESTRICTION_CURRENT_VALUE, // put in the new value of the current usage we calculated + the number of - // tokens we just created + // tokens we just created, plus any cache tokens at their configured weight ((Number) userRestrictionMap.get(AbstractModelEngineResponse.USAGE_RESTRICTION_CURRENT_VALUE)) .intValue() + modelResponse.getNumberOfTokensInPrompt() - + modelResponse.getNumberOfTokensInResponse()); + + modelResponse.getNumberOfTokensInResponse() + (int) Math.round(weightedCacheTokens)); } else if (Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE.equals(restrictionMode)) { diff --git a/src/prerna/engine/impl/model/inferencetracking/ModelInferenceLogsUtils.java b/src/prerna/engine/impl/model/inferencetracking/ModelInferenceLogsUtils.java index 2e725fecebc..87ad3218cc7 100644 --- a/src/prerna/engine/impl/model/inferencetracking/ModelInferenceLogsUtils.java +++ b/src/prerna/engine/impl/model/inferencetracking/ModelInferenceLogsUtils.java @@ -2011,10 +2011,17 @@ private static void deleteFeedbackEntry(String messageId) { * @param currentDateTime reference date/time * @param frequency window frequency ({@code DAY}, {@code WEEK}, * {@code MONTH}, {@code YEAR}, {@code ALL_TIME}) + * @param cacheReadWeightPercent percentage of a normal token a cache read + * token counts as (100 = same as a normal + * token); ignored outside token mode + * @param cacheWriteWeightPercent percentage of a normal token a cache + * creation token counts as; ignored outside + * token mode * @return aggregate usage value, or {@code null} if unavailable */ public static Number getTotalTokensOrTotalResponseTime(String restrictionMode, User user, String engineId, - ZonedDateTime currentDateTime, String frequency) { + ZonedDateTime currentDateTime, String frequency, double cacheReadWeightPercent, + double cacheWriteWeightPercent) { IRDBMSEngine modelInferenceLogsDb = SystemEngineRegistry.getModelInferenceLogsDb(); if (restrictionMode == null) { throw new IllegalArgumentException("Must pass in a valid restriction mode"); @@ -2029,16 +2036,20 @@ public static Number getTotalTokensOrTotalResponseTime(String restrictionMode, U ZonedDateTime startDate = dates.get("start"); ZonedDateTime endDate = dates.get("end"); + boolean isTokenMode = restrictionMode.equalsIgnoreCase(Constants.MODEL_TOKEN_RESTRICTION_VALUE) + || restrictionMode.equalsIgnoreCase(Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE); String sumColumn = null; - if (restrictionMode.equalsIgnoreCase(Constants.MODEL_TOKEN_RESTRICTION_VALUE)) { - sumColumn = " SUM(MESSAGE_TOKENS) "; + if (isTokenMode) { + // cache tokens are summed separately so the caller can weight them; they are + // not already folded into MESSAGE_TOKENS + sumColumn = " SUM(MESSAGE_TOKENS), SUM(CACHE_READ_TOKENS), SUM(CACHE_CREATION_TOKENS) "; } else if (restrictionMode.equalsIgnoreCase(Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE)) { sumColumn = " SUM(RESPONSE_TIME) "; } // SQL query to fetch the total tokens or response time String query = "SELECT " + sumColumn - + " AS \"current_usage\" FROM MESSAGE WHERE USER_ID=? AND AGENT_ID=? AND DATE_CREATED BETWEEN ? AND ?"; + + " FROM MESSAGE WHERE USER_ID=? AND AGENT_ID=? AND DATE_CREATED BETWEEN ? AND ?"; PreparedStatement ps = null; ResultSet rs = null; try { @@ -2053,14 +2064,15 @@ public static Number getTotalTokensOrTotalResponseTime(String restrictionMode, U ps.getConnection(), ps, query, false); if (wrapper.hasNext()) { - Number retNum = (Number) wrapper.next().getValues()[0]; + Object[] values = wrapper.next().getValues(); + if (isTokenMode) { + return weightedTokenTotal(values, cacheReadWeightPercent, cacheWriteWeightPercent); + } // if this is null // that means there are no logs currently for this model // we will treat this as 0 usage - if (retNum == null) { - return 0; - } - return retNum; + Number retNum = (Number) values[0]; + return retNum == null ? 0 : retNum; } } catch (Exception e) { classLogger.error( @@ -2072,6 +2084,30 @@ public static Number getTotalTokensOrTotalResponseTime(String restrictionMode, U return null; } + /** + * Combine the base token sum with weighted cache tokens. Cache reads and cache + * writes are billed (and therefore should count toward a limit) differently + * from a fresh token, so each gets its own admin-configured percentage rather + * than counting 1:1 by default... a weight of 100 reproduces the un-weighted + * total. + * + * @param values positional result row: [messageTokens, + * cacheReadTokens, cacheCreationTokens] + * @param cacheReadWeightPercent percentage of a normal token a cache read + * token counts as + * @param cacheWriteWeightPercent percentage of a normal token a cache + * creation token counts as + * @return the weighted total, as a double + */ + private static Number weightedTokenTotal(Object[] values, double cacheReadWeightPercent, + double cacheWriteWeightPercent) { + double messageTokens = values[0] == null ? 0 : ((Number) values[0]).doubleValue(); + double cacheReadTokens = values[1] == null ? 0 : ((Number) values[1]).doubleValue(); + double cacheCreationTokens = values[2] == null ? 0 : ((Number) values[2]).doubleValue(); + return messageTokens + cacheReadTokens * (cacheReadWeightPercent / 100.0) + + cacheCreationTokens * (cacheWriteWeightPercent / 100.0); + } + /** * Calculates total token usage or response time for a user across engines while * excluding restricted engines. @@ -2083,10 +2119,17 @@ public static Number getTotalTokensOrTotalResponseTime(String restrictionMode, U * @param currentDateTime reference date/time * @param frequency window frequency ({@code DAY}, {@code WEEK}, * {@code MONTH}, {@code YEAR}, {@code ALL_TIME}) + * @param cacheReadWeightPercent percentage of a normal token a cache read + * token counts as (100 = same as a normal + * token); ignored outside token mode + * @param cacheWriteWeightPercent percentage of a normal token a cache + * creation token counts as; ignored outside + * token mode * @return aggregate usage value, or {@code null} if unavailable */ public static Number getTotalUsageForUser(String restrictionMode, User user, String engineId, - ZonedDateTime currentDateTime, String frequency) { + ZonedDateTime currentDateTime, String frequency, double cacheReadWeightPercent, + double cacheWriteWeightPercent) { IRDBMSEngine modelInferenceLogsDb = SystemEngineRegistry.getModelInferenceLogsDb(); if (restrictionMode == null) { throw new IllegalArgumentException("Must pass in a valid restriction mode"); @@ -2118,17 +2161,20 @@ public static Number getTotalUsageForUser(String restrictionMode, User user, Str // Step 3: Determine which column to sum (tokens or response time) based on // restrictionMode + boolean isTokenMode = restrictionMode.equalsIgnoreCase(Constants.MODEL_TOKEN_RESTRICTION_VALUE) + || restrictionMode.equalsIgnoreCase(Constants.MODEL_TOKEN_CACHE_RESTRICTION_VALUE); String sumColumn = null; - if (restrictionMode.equalsIgnoreCase(Constants.MODEL_TOKEN_RESTRICTION_VALUE)) { - sumColumn = " SUM(MESSAGE_TOKENS) "; + if (isTokenMode) { + // cache tokens are summed separately so the caller can weight them; they are + // not already folded into MESSAGE_TOKENS + sumColumn = " SUM(MESSAGE_TOKENS), SUM(CACHE_READ_TOKENS), SUM(CACHE_CREATION_TOKENS) "; } else if (restrictionMode.equalsIgnoreCase(Constants.MODEL_COMPUTE_TIME_RESTRICTION_VALUE)) { sumColumn = " SUM(RESPONSE_TIME) "; } // Step 4: Get total usage for the user excluding the engines in the // engineIdList - String query = "SELECT " + sumColumn - + " AS \"current_usage\" FROM MESSAGE WHERE USER_ID=? AND DATE_CREATED BETWEEN ? AND ? " + String query = "SELECT " + sumColumn + " FROM MESSAGE WHERE USER_ID=? AND DATE_CREATED BETWEEN ? AND ? " + excludePSString; PreparedStatement ps = null; ResultSet rs = null; @@ -2149,14 +2195,15 @@ public static Number getTotalUsageForUser(String restrictionMode, User user, Str ps.getConnection(), ps, query, false); if (wrapper.hasNext()) { - Number retNum = (Number) wrapper.next().getValues()[0]; + Object[] values = wrapper.next().getValues(); + if (isTokenMode) { + return weightedTokenTotal(values, cacheReadWeightPercent, cacheWriteWeightPercent); + } // if this is null // that means there are no logs currently for this model // we will treat this as 0 usage - if (retNum == null) { - return 0; - } - return retNum; + Number retNum = (Number) values[0]; + return retNum == null ? 0 : retNum; } } catch (Exception e) { classLogger.error("Failed to calculate total usage for userId '{}', restrictionMode '{}', frequency '{}'.", diff --git a/src/prerna/engine/impl/model/responses/AbstractModelEngineResponse.java b/src/prerna/engine/impl/model/responses/AbstractModelEngineResponse.java index 27af639d961..0899b857a22 100644 --- a/src/prerna/engine/impl/model/responses/AbstractModelEngineResponse.java +++ b/src/prerna/engine/impl/model/responses/AbstractModelEngineResponse.java @@ -48,6 +48,8 @@ public abstract class AbstractModelEngineResponse implements Serializable { public static final String USAGE_RESTRICTION_MODE = "restrictedBy"; public static final String USAGE_RESTRICTION_CURRENT_VALUE = "currentValue"; public static final String USAGE_RESTRICTION_MAX_VALUE = "maxValue"; + public static final String USAGE_RESTRICTION_CACHE_READ_WEIGHT = "cacheReadWeight"; + public static final String USAGE_RESTRICTION_CACHE_WRITE_WEIGHT = "cacheWriteWeight"; protected T response; protected Integer numberOfTokensInPrompt; diff --git a/src/prerna/util/Constants.java b/src/prerna/util/Constants.java index 7ac02183028..3b3f3dda401 100644 --- a/src/prerna/util/Constants.java +++ b/src/prerna/util/Constants.java @@ -500,6 +500,10 @@ public class Constants { public static final String SUPPORTED_PARAMETERS = "SUPPORTED_PARAMETERS"; public static final String REASONING_CONFIG = "REASONING_CONFIG"; public static final String BENCHMARKS = "BENCHMARKS"; + // Percentage (0-1000, default 100 when unset) applied to cache read/write + // tokens when they count toward a member's "token" usage restriction. + public static final String CACHE_READ_WEIGHT = "CACHE_READ_WEIGHT"; + public static final String CACHE_WRITE_WEIGHT = "CACHE_WRITE_WEIGHT"; // Compare Databases public static final String NEW_DB_COMBOBOX = "newDBComboBox"; @@ -1023,6 +1027,10 @@ public class Constants { // model restriction types public static final String MODEL_TOKEN_RESTRICTION_VALUE = "token"; public static final String MODEL_COMPUTE_TIME_RESTRICTION_VALUE = "compute"; + // Same as MODEL_TOKEN_RESTRICTION_VALUE, but also counts cache read/write + // tokens toward the limit at the model's configured CACHE_READ_WEIGHT / + // CACHE_WRITE_WEIGHT (see SecurityModelMetadataUtils). + public static final String MODEL_TOKEN_CACHE_RESTRICTION_VALUE = "token_cache"; // External Permission Management public static final String EXTERNAL_PERMISSION_MANAGEMENT_ENABLED = "EXTERNAL_PERMISSION_MANAGEMENT_ENABLED";