diff --git a/src/prerna/auth/utils/SecurityEngineUtils.java b/src/prerna/auth/utils/SecurityEngineUtils.java index 16f7e20a06..06471715d5 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 e2016b650a..50a91dca31 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 5c33583f1f..45070e4cf9 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 8f77716099..5ae83e40ba 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 2e725feceb..87ad3218cc 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 27af639d96..0899b857a2 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 7ac0218302..3b3f3dda40 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";