floe-guard 0.1.0 → 0.2.1
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +32 -4
- package/dist/index.cjs +328 -53
- package/dist/index.cjs.map +1 -1
- package/dist/index.d.cts +131 -15
- package/dist/index.d.ts +131 -15
- package/dist/index.js +328 -53
- package/dist/index.js.map +1 -1
- package/package.json +4 -3
- package/src/cost_map.json +90 -0
package/dist/index.js
CHANGED
|
@@ -80,6 +80,12 @@ var cost_map_default = {
|
|
|
80
80
|
litellm_provider: "anthropic",
|
|
81
81
|
mode: "chat"
|
|
82
82
|
},
|
|
83
|
+
"claude-fable-5": {
|
|
84
|
+
input_cost_per_token: 1e-5,
|
|
85
|
+
output_cost_per_token: 5e-5,
|
|
86
|
+
litellm_provider: "anthropic",
|
|
87
|
+
mode: "chat"
|
|
88
|
+
},
|
|
83
89
|
"claude-haiku-4-5": {
|
|
84
90
|
input_cost_per_token: 1e-6,
|
|
85
91
|
output_cost_per_token: 5e-6,
|
|
@@ -176,6 +182,12 @@ var cost_map_default = {
|
|
|
176
182
|
litellm_provider: "anthropic",
|
|
177
183
|
mode: "chat"
|
|
178
184
|
},
|
|
185
|
+
"claude-sonnet-5": {
|
|
186
|
+
input_cost_per_token: 2e-6,
|
|
187
|
+
output_cost_per_token: 1e-5,
|
|
188
|
+
litellm_provider: "anthropic",
|
|
189
|
+
mode: "chat"
|
|
190
|
+
},
|
|
179
191
|
"ft:gpt-3.5-turbo": {
|
|
180
192
|
input_cost_per_token: 3e-6,
|
|
181
193
|
output_cost_per_token: 6e-6,
|
|
@@ -626,6 +638,30 @@ var cost_map_default = {
|
|
|
626
638
|
litellm_provider: "openai",
|
|
627
639
|
mode: "chat"
|
|
628
640
|
},
|
|
641
|
+
"gpt-5.6": {
|
|
642
|
+
input_cost_per_token: 5e-6,
|
|
643
|
+
output_cost_per_token: 3e-5,
|
|
644
|
+
litellm_provider: "openai",
|
|
645
|
+
mode: "chat"
|
|
646
|
+
},
|
|
647
|
+
"gpt-5.6-luna": {
|
|
648
|
+
input_cost_per_token: 1e-6,
|
|
649
|
+
output_cost_per_token: 6e-6,
|
|
650
|
+
litellm_provider: "openai",
|
|
651
|
+
mode: "chat"
|
|
652
|
+
},
|
|
653
|
+
"gpt-5.6-sol": {
|
|
654
|
+
input_cost_per_token: 5e-6,
|
|
655
|
+
output_cost_per_token: 3e-5,
|
|
656
|
+
litellm_provider: "openai",
|
|
657
|
+
mode: "chat"
|
|
658
|
+
},
|
|
659
|
+
"gpt-5.6-terra": {
|
|
660
|
+
input_cost_per_token: 25e-7,
|
|
661
|
+
output_cost_per_token: 15e-6,
|
|
662
|
+
litellm_provider: "openai",
|
|
663
|
+
mode: "chat"
|
|
664
|
+
},
|
|
629
665
|
"gpt-audio": {
|
|
630
666
|
input_cost_per_token: 25e-7,
|
|
631
667
|
output_cost_per_token: 1e-5,
|
|
@@ -680,6 +716,18 @@ var cost_map_default = {
|
|
|
680
716
|
litellm_provider: "openai",
|
|
681
717
|
mode: "chat"
|
|
682
718
|
},
|
|
719
|
+
"gpt-realtime-2.1": {
|
|
720
|
+
input_cost_per_token: 4e-6,
|
|
721
|
+
output_cost_per_token: 24e-6,
|
|
722
|
+
litellm_provider: "openai",
|
|
723
|
+
mode: "chat"
|
|
724
|
+
},
|
|
725
|
+
"gpt-realtime-2.1-mini": {
|
|
726
|
+
input_cost_per_token: 6e-7,
|
|
727
|
+
output_cost_per_token: 24e-7,
|
|
728
|
+
litellm_provider: "openai",
|
|
729
|
+
mode: "chat"
|
|
730
|
+
},
|
|
683
731
|
"gpt-realtime-2025-08-28": {
|
|
684
732
|
input_cost_per_token: 4e-6,
|
|
685
733
|
output_cost_per_token: 16e-6,
|
|
@@ -704,6 +752,24 @@ var cost_map_default = {
|
|
|
704
752
|
litellm_provider: "openai",
|
|
705
753
|
mode: "chat"
|
|
706
754
|
},
|
|
755
|
+
"llama-3.1-8b-instant": {
|
|
756
|
+
input_cost_per_token: 5e-8,
|
|
757
|
+
output_cost_per_token: 8e-8,
|
|
758
|
+
litellm_provider: "groq",
|
|
759
|
+
mode: "chat"
|
|
760
|
+
},
|
|
761
|
+
"llama-3.3-70b-versatile": {
|
|
762
|
+
input_cost_per_token: 59e-8,
|
|
763
|
+
output_cost_per_token: 79e-8,
|
|
764
|
+
litellm_provider: "groq",
|
|
765
|
+
mode: "chat"
|
|
766
|
+
},
|
|
767
|
+
"meta-llama/llama-4-scout-17b-16e-instruct": {
|
|
768
|
+
input_cost_per_token: 11e-8,
|
|
769
|
+
output_cost_per_token: 34e-8,
|
|
770
|
+
litellm_provider: "groq",
|
|
771
|
+
mode: "chat"
|
|
772
|
+
},
|
|
707
773
|
o1: {
|
|
708
774
|
input_cost_per_token: 15e-6,
|
|
709
775
|
output_cost_per_token: 6e-5,
|
|
@@ -752,6 +818,30 @@ var cost_map_default = {
|
|
|
752
818
|
litellm_provider: "openai",
|
|
753
819
|
mode: "chat"
|
|
754
820
|
},
|
|
821
|
+
"openai/gpt-oss-120b": {
|
|
822
|
+
input_cost_per_token: 15e-8,
|
|
823
|
+
output_cost_per_token: 6e-7,
|
|
824
|
+
litellm_provider: "groq",
|
|
825
|
+
mode: "chat"
|
|
826
|
+
},
|
|
827
|
+
"openai/gpt-oss-20b": {
|
|
828
|
+
input_cost_per_token: 75e-9,
|
|
829
|
+
output_cost_per_token: 3e-7,
|
|
830
|
+
litellm_provider: "groq",
|
|
831
|
+
mode: "chat"
|
|
832
|
+
},
|
|
833
|
+
"openai/gpt-oss-safeguard-20b": {
|
|
834
|
+
input_cost_per_token: 75e-9,
|
|
835
|
+
output_cost_per_token: 3e-7,
|
|
836
|
+
litellm_provider: "groq",
|
|
837
|
+
mode: "chat"
|
|
838
|
+
},
|
|
839
|
+
"qwen/qwen3-32b": {
|
|
840
|
+
input_cost_per_token: 29e-8,
|
|
841
|
+
output_cost_per_token: 59e-8,
|
|
842
|
+
litellm_provider: "groq",
|
|
843
|
+
mode: "chat"
|
|
844
|
+
},
|
|
755
845
|
"text-embedding-3-large": {
|
|
756
846
|
input_cost_per_token: 13e-8,
|
|
757
847
|
output_cost_per_token: 0,
|
|
@@ -780,39 +870,64 @@ var cost_map_default = {
|
|
|
780
870
|
|
|
781
871
|
// src/pricing.ts
|
|
782
872
|
var COST_MAP = cost_map_default;
|
|
783
|
-
|
|
873
|
+
var PROVIDER_PREFIXES = /* @__PURE__ */ new Set(["groq"]);
|
|
874
|
+
var DATE_SUFFIX = /-(?:\d{8}|\d{4}-\d{2}-\d{2})$/;
|
|
875
|
+
function candidateGroups(model) {
|
|
784
876
|
const m = model.trim();
|
|
785
|
-
const
|
|
786
|
-
|
|
877
|
+
const base = [m];
|
|
878
|
+
const firstSlash = m.indexOf("/");
|
|
879
|
+
if (firstSlash !== -1 && PROVIDER_PREFIXES.has(m.slice(0, firstSlash))) {
|
|
880
|
+
base.push(m.slice(firstSlash + 1));
|
|
881
|
+
}
|
|
882
|
+
const lastSlash = m.lastIndexOf("/");
|
|
883
|
+
if (lastSlash !== -1) {
|
|
884
|
+
base.push(m.slice(lastSlash + 1));
|
|
885
|
+
}
|
|
886
|
+
const exact = [];
|
|
887
|
+
for (const cand of base) {
|
|
888
|
+
if (cand && !exact.includes(cand)) exact.push(cand);
|
|
889
|
+
}
|
|
890
|
+
const stripped = [];
|
|
891
|
+
for (const cand of exact) {
|
|
892
|
+
const c = cand.replace(DATE_SUFFIX, "");
|
|
893
|
+
if (c && !exact.includes(c) && !stripped.includes(c)) stripped.push(c);
|
|
894
|
+
}
|
|
895
|
+
return [exact, stripped];
|
|
787
896
|
}
|
|
788
897
|
function bothFinite(a, b) {
|
|
789
898
|
return typeof a === "number" && typeof b === "number" && Number.isFinite(a) && Number.isFinite(b);
|
|
790
899
|
}
|
|
791
900
|
function resolvePrice(model, overrides) {
|
|
792
|
-
const
|
|
793
|
-
|
|
794
|
-
|
|
795
|
-
|
|
796
|
-
|
|
797
|
-
|
|
798
|
-
|
|
799
|
-
|
|
800
|
-
|
|
801
|
-
|
|
901
|
+
for (const cands of candidateGroups(model)) {
|
|
902
|
+
if (overrides) {
|
|
903
|
+
for (const cand of cands) {
|
|
904
|
+
const ov = Object.prototype.hasOwnProperty.call(overrides, cand) ? overrides[cand] : void 0;
|
|
905
|
+
if (ov !== void 0) {
|
|
906
|
+
if (bothFinite(ov.inputCostPerToken, ov.outputCostPerToken)) {
|
|
907
|
+
return {
|
|
908
|
+
inputCostPerToken: ov.inputCostPerToken,
|
|
909
|
+
outputCostPerToken: ov.outputCostPerToken,
|
|
910
|
+
source: "override"
|
|
911
|
+
};
|
|
912
|
+
}
|
|
913
|
+
return null;
|
|
914
|
+
}
|
|
802
915
|
}
|
|
803
|
-
|
|
916
|
+
}
|
|
917
|
+
for (const cand of cands) {
|
|
918
|
+
const entry = Object.prototype.hasOwnProperty.call(COST_MAP, cand) ? COST_MAP[cand] : void 0;
|
|
919
|
+
if (!entry) continue;
|
|
920
|
+
const input = entry.input_cost_per_token;
|
|
921
|
+
const output = entry.output_cost_per_token;
|
|
922
|
+
if (!bothFinite(input, output)) return null;
|
|
923
|
+
return {
|
|
924
|
+
inputCostPerToken: input,
|
|
925
|
+
outputCostPerToken: output,
|
|
926
|
+
source: "cost_map"
|
|
927
|
+
};
|
|
804
928
|
}
|
|
805
929
|
}
|
|
806
|
-
|
|
807
|
-
if (!entry) return null;
|
|
808
|
-
const input = entry.input_cost_per_token;
|
|
809
|
-
const output = entry.output_cost_per_token;
|
|
810
|
-
if (!bothFinite(input, output)) return null;
|
|
811
|
-
return {
|
|
812
|
-
inputCostPerToken: input,
|
|
813
|
-
outputCostPerToken: output,
|
|
814
|
-
source: "cost_map"
|
|
815
|
-
};
|
|
930
|
+
return null;
|
|
816
931
|
}
|
|
817
932
|
function priceTokens(priced, promptTokens, completionTokens) {
|
|
818
933
|
const p = Math.max(0, promptTokens);
|
|
@@ -831,9 +946,12 @@ var BudgetGuard = class {
|
|
|
831
946
|
spentUsd = 0;
|
|
832
947
|
priceOverrides;
|
|
833
948
|
failClosed;
|
|
949
|
+
nearLimitBps;
|
|
834
950
|
onBlock;
|
|
835
951
|
/** Cost of the most recent priced call, used to predict the next one. */
|
|
836
952
|
lastCost = 0;
|
|
953
|
+
/** USD held for in-flight calls (reserved, not yet settled). Counts toward the ceiling. */
|
|
954
|
+
reserved = 0;
|
|
837
955
|
/**
|
|
838
956
|
* @param limitUsd the spend ceiling, in USD. `0` blocks the very first call.
|
|
839
957
|
*/
|
|
@@ -843,35 +961,81 @@ var BudgetGuard = class {
|
|
|
843
961
|
`limitUsd must be a finite, non-negative number, got ${limitUsd}`
|
|
844
962
|
);
|
|
845
963
|
}
|
|
964
|
+
const nearLimitBps = options.nearLimitBps === void 0 ? 8e3 : options.nearLimitBps;
|
|
965
|
+
if (!Number.isInteger(nearLimitBps) || nearLimitBps < 0 || nearLimitBps > 1e4) {
|
|
966
|
+
throw new RangeError(
|
|
967
|
+
`nearLimitBps must be an integer in 0..10000, got ${nearLimitBps}`
|
|
968
|
+
);
|
|
969
|
+
}
|
|
846
970
|
this.limitUsd = limitUsd;
|
|
847
971
|
this.priceOverrides = options.priceOverrides;
|
|
848
972
|
this.failClosed = options.failClosed ?? true;
|
|
849
973
|
this.onBlock = options.onBlock ?? defaultOnBlock;
|
|
974
|
+
this.nearLimitBps = nearLimitBps;
|
|
850
975
|
}
|
|
851
976
|
/**
|
|
852
977
|
* Throw {@link BudgetExceeded} if the next call would cross the ceiling.
|
|
853
978
|
*
|
|
854
979
|
* Call this immediately before each LLM request. The "next call" is estimated
|
|
855
980
|
* from the last recorded call's cost (override with `estimatedNextCost`); the
|
|
856
|
-
* first call is always allowed unless the ceiling is already met.
|
|
857
|
-
*
|
|
981
|
+
* first call is always allowed unless the ceiling is already met. In-flight
|
|
982
|
+
* reservations count toward the total, so this stays correct alongside
|
|
983
|
+
* {@link BudgetGuard.reserve}.
|
|
984
|
+
*
|
|
985
|
+
* Note: `check` is a non-binding peek. For parallel calls, use `reserve()` /
|
|
986
|
+
* `settle()`, which hold the estimate across the await.
|
|
858
987
|
*/
|
|
859
988
|
check(estimatedNextCost) {
|
|
860
|
-
const
|
|
861
|
-
|
|
862
|
-
|
|
989
|
+
const rawEstimate = estimatedNextCost === void 0 ? this.lastCost : estimatedNextCost;
|
|
990
|
+
if (!Number.isFinite(rawEstimate)) {
|
|
991
|
+
throw new RangeError(
|
|
992
|
+
`estimatedNextCost must be a finite number, got ${rawEstimate}`
|
|
993
|
+
);
|
|
994
|
+
}
|
|
995
|
+
const estimate = Math.max(0, rawEstimate);
|
|
996
|
+
const committed = this.spentUsd + this.reserved;
|
|
997
|
+
if (committed > this.limitUsd - EPS || committed + estimate > this.limitUsd + EPS) {
|
|
863
998
|
this.onBlock(this.spentUsd, this.limitUsd);
|
|
864
999
|
throw new BudgetExceeded(this.spentUsd, this.limitUsd);
|
|
865
1000
|
}
|
|
866
1001
|
}
|
|
867
1002
|
/**
|
|
868
|
-
*
|
|
1003
|
+
* Atomically check the ceiling AND hold the estimated cost in flight.
|
|
869
1004
|
*
|
|
870
|
-
*
|
|
871
|
-
*
|
|
872
|
-
*
|
|
1005
|
+
* The concurrency-safe enforcement path: call before the request and hold the
|
|
1006
|
+
* returned reservation across the await, so parallel callers can't all clear
|
|
1007
|
+
* the same stale total. Throws {@link BudgetExceeded} (without reserving) if
|
|
1008
|
+
* the reservation would cross the ceiling. Returns the reservation handle to
|
|
1009
|
+
* pass to {@link BudgetGuard.settle} (or {@link BudgetGuard.release} on error).
|
|
1010
|
+
* `estimatedCost` defaults to the last call's cost.
|
|
873
1011
|
*/
|
|
874
|
-
|
|
1012
|
+
reserve(estimatedCost) {
|
|
1013
|
+
const rawEstimate = estimatedCost === void 0 ? this.lastCost : estimatedCost;
|
|
1014
|
+
if (!Number.isFinite(rawEstimate)) {
|
|
1015
|
+
throw new RangeError(
|
|
1016
|
+
`estimatedCost must be a finite number, got ${rawEstimate}`
|
|
1017
|
+
);
|
|
1018
|
+
}
|
|
1019
|
+
const estimate = Math.max(0, rawEstimate);
|
|
1020
|
+
const committed = this.spentUsd + this.reserved;
|
|
1021
|
+
if (committed > this.limitUsd - EPS || committed + estimate > this.limitUsd + EPS) {
|
|
1022
|
+
this.onBlock(this.spentUsd, this.limitUsd);
|
|
1023
|
+
throw new BudgetExceeded(this.spentUsd, this.limitUsd);
|
|
1024
|
+
}
|
|
1025
|
+
this.reserved += estimate;
|
|
1026
|
+
return estimate;
|
|
1027
|
+
}
|
|
1028
|
+
/**
|
|
1029
|
+
* Release a reservation and record the actual cost. `record` is `settle` with
|
|
1030
|
+
* no reservation. Returns the USD cost of this call; unpriceable-model handling
|
|
1031
|
+
* matches {@link BudgetGuard.record}, and any held reservation is released even
|
|
1032
|
+
* on the warn-and-skip path.
|
|
1033
|
+
*/
|
|
1034
|
+
settle(model, promptTokens, completionTokens, options = {}) {
|
|
1035
|
+
const reserved = options.reserved ?? 0;
|
|
1036
|
+
if (!Number.isFinite(reserved) || reserved < 0) {
|
|
1037
|
+
throw new RangeError(`reserved must be a finite, non-negative number, got ${reserved}`);
|
|
1038
|
+
}
|
|
875
1039
|
let overrides = this.priceOverrides;
|
|
876
1040
|
if (options.price !== void 0) {
|
|
877
1041
|
overrides = { ...overrides ?? {}, [model]: options.price };
|
|
@@ -881,12 +1045,22 @@ var BudgetGuard = class {
|
|
|
881
1045
|
console.warn(
|
|
882
1046
|
`Cannot price model '${model}': not in the bundled cost map and no manual price given. The budget guard cannot enforce a ceiling on spend it cannot measure \u2014 pass { price } or set it in priceOverrides.`
|
|
883
1047
|
);
|
|
1048
|
+
this.release(reserved);
|
|
884
1049
|
if (this.failClosed) {
|
|
885
1050
|
throw new UnpriceableModelError(model);
|
|
886
1051
|
}
|
|
887
1052
|
return 0;
|
|
888
1053
|
}
|
|
889
|
-
|
|
1054
|
+
let cost;
|
|
1055
|
+
try {
|
|
1056
|
+
cost = priceTokens(priced, promptTokens, completionTokens);
|
|
1057
|
+
} catch (err) {
|
|
1058
|
+
this.release(reserved);
|
|
1059
|
+
throw err;
|
|
1060
|
+
}
|
|
1061
|
+
if (reserved) {
|
|
1062
|
+
this.reserved = Math.max(0, this.reserved - reserved);
|
|
1063
|
+
}
|
|
890
1064
|
this.spentUsd += cost;
|
|
891
1065
|
if (this.spentUsd - this.limitUsd > 0 && this.spentUsd - this.limitUsd < EPS) {
|
|
892
1066
|
this.spentUsd = this.limitUsd;
|
|
@@ -894,9 +1068,55 @@ var BudgetGuard = class {
|
|
|
894
1068
|
this.lastCost = cost;
|
|
895
1069
|
return cost;
|
|
896
1070
|
}
|
|
897
|
-
/**
|
|
1071
|
+
/**
|
|
1072
|
+
* Price one response's tokens offline and add the cost to the total.
|
|
1073
|
+
*
|
|
1074
|
+
* Returns the USD cost of this call. If the model is unpriceable and no `price`
|
|
1075
|
+
* is given, behaviour depends on `failClosed`: warn + throw (default), or
|
|
1076
|
+
* warn + skip accrual.
|
|
1077
|
+
*/
|
|
1078
|
+
record(model, promptTokens, completionTokens, options = {}) {
|
|
1079
|
+
return this.settle(model, promptTokens, completionTokens, {
|
|
1080
|
+
reserved: 0,
|
|
1081
|
+
price: options.price
|
|
1082
|
+
});
|
|
1083
|
+
}
|
|
1084
|
+
/**
|
|
1085
|
+
* Drop an in-flight reservation without recording spend (e.g. the call failed
|
|
1086
|
+
* before producing usage). Safe to call with `0`.
|
|
1087
|
+
*/
|
|
1088
|
+
release(reserved) {
|
|
1089
|
+
if (!Number.isFinite(reserved) || reserved < 0) {
|
|
1090
|
+
throw new RangeError(`reserved must be a finite, non-negative number, got ${reserved}`);
|
|
1091
|
+
}
|
|
1092
|
+
if (!reserved) return;
|
|
1093
|
+
this.reserved = Math.max(0, this.reserved - reserved);
|
|
1094
|
+
}
|
|
1095
|
+
/** USD left before the ceiling, net of in-flight reservations (never negative). */
|
|
898
1096
|
get remainingUsd() {
|
|
899
|
-
return Math.max(0, this.limitUsd - this.spentUsd);
|
|
1097
|
+
return Math.max(0, this.limitUsd - this.spentUsd - this.reserved);
|
|
1098
|
+
}
|
|
1099
|
+
/**
|
|
1100
|
+
* Context-aware spend advisory for this budget — see {@link BudgetAdvisory}.
|
|
1101
|
+
*
|
|
1102
|
+
* `nearLimit` flips once utilization reaches `nearLimitBps` (default 80%), so an
|
|
1103
|
+
* agent can taper *before* the hard-stop. Advisory only: read it to adapt;
|
|
1104
|
+
* {@link BudgetGuard.check} is what enforces the ceiling.
|
|
1105
|
+
*/
|
|
1106
|
+
advisory() {
|
|
1107
|
+
const usedBps = this.limitUsd <= 0 ? 1e4 : Math.max(0, Math.min(1e4, Math.floor(this.spentUsd / this.limitUsd * 1e4 + 1e-9)));
|
|
1108
|
+
return {
|
|
1109
|
+
nearLimit: usedBps >= this.nearLimitBps,
|
|
1110
|
+
usedBps,
|
|
1111
|
+
// Settled budget: limit minus accrued spend, deliberately NOT net of
|
|
1112
|
+
// in-flight reservations. Unlike the remainingUsd getter (which subtracts
|
|
1113
|
+
// `reserved`), the advisory is a soft utilization signal about money already
|
|
1114
|
+
// spent, while the getter reports what a new call can still claim.
|
|
1115
|
+
remainingUsd: Math.max(0, this.limitUsd - this.spentUsd),
|
|
1116
|
+
limitUsd: this.limitUsd,
|
|
1117
|
+
spentUsd: this.spentUsd,
|
|
1118
|
+
scope: "local"
|
|
1119
|
+
};
|
|
900
1120
|
}
|
|
901
1121
|
};
|
|
902
1122
|
function defaultOnBlock(spentUsd, limitUsd) {
|
|
@@ -908,33 +1128,88 @@ function defaultOnBlock(spentUsd, limitUsd) {
|
|
|
908
1128
|
}
|
|
909
1129
|
|
|
910
1130
|
// src/middleware.ts
|
|
1131
|
+
function usageTokens(modelId, usage) {
|
|
1132
|
+
const u = usage;
|
|
1133
|
+
const promptTokens = u?.promptTokens ?? u?.inputTokens;
|
|
1134
|
+
const completionTokens = u?.completionTokens ?? u?.outputTokens;
|
|
1135
|
+
if (typeof promptTokens !== "number" || typeof completionTokens !== "number") {
|
|
1136
|
+
throw new Error(
|
|
1137
|
+
`Model '${modelId}' reported no token usage \u2014 the budget guard cannot meter spend it cannot see, so this call is rejected rather than treated as free.`
|
|
1138
|
+
);
|
|
1139
|
+
}
|
|
1140
|
+
return { promptTokens, completionTokens };
|
|
1141
|
+
}
|
|
911
1142
|
function budgetGuardMiddleware(guard) {
|
|
912
1143
|
return {
|
|
913
1144
|
async wrapGenerate({ doGenerate, model }) {
|
|
914
|
-
guard.
|
|
915
|
-
|
|
916
|
-
|
|
917
|
-
|
|
918
|
-
|
|
919
|
-
|
|
920
|
-
|
|
921
|
-
|
|
1145
|
+
const reserved = guard.reserve();
|
|
1146
|
+
let result;
|
|
1147
|
+
try {
|
|
1148
|
+
result = await doGenerate();
|
|
1149
|
+
} catch (err) {
|
|
1150
|
+
guard.release(reserved);
|
|
1151
|
+
throw err;
|
|
1152
|
+
}
|
|
1153
|
+
let handled = false;
|
|
1154
|
+
try {
|
|
1155
|
+
const { promptTokens, completionTokens } = usageTokens(
|
|
1156
|
+
model.modelId,
|
|
1157
|
+
result?.usage
|
|
1158
|
+
);
|
|
1159
|
+
handled = true;
|
|
1160
|
+
guard.settle(model.modelId, promptTokens, completionTokens, { reserved });
|
|
1161
|
+
return result;
|
|
1162
|
+
} catch (err) {
|
|
1163
|
+
if (!handled) guard.release(reserved);
|
|
1164
|
+
throw err;
|
|
1165
|
+
}
|
|
922
1166
|
},
|
|
923
1167
|
async wrapStream({ doStream, model }) {
|
|
924
|
-
guard.
|
|
925
|
-
|
|
1168
|
+
const reserved = guard.reserve();
|
|
1169
|
+
let streamResult;
|
|
1170
|
+
try {
|
|
1171
|
+
streamResult = await doStream();
|
|
1172
|
+
} catch (err) {
|
|
1173
|
+
guard.release(reserved);
|
|
1174
|
+
throw err;
|
|
1175
|
+
}
|
|
1176
|
+
const { stream, ...rest } = streamResult;
|
|
1177
|
+
let handled = false;
|
|
926
1178
|
const guarded = stream.pipeThrough(
|
|
927
1179
|
new TransformStream({
|
|
928
1180
|
transform(chunk, controller) {
|
|
929
|
-
if (chunk
|
|
930
|
-
|
|
931
|
-
|
|
932
|
-
|
|
933
|
-
|
|
934
|
-
|
|
1181
|
+
if (chunk?.type === "finish" && !handled) {
|
|
1182
|
+
try {
|
|
1183
|
+
const { promptTokens, completionTokens } = usageTokens(
|
|
1184
|
+
model.modelId,
|
|
1185
|
+
chunk.usage
|
|
1186
|
+
);
|
|
1187
|
+
handled = true;
|
|
1188
|
+
guard.settle(model.modelId, promptTokens, completionTokens, { reserved });
|
|
1189
|
+
} catch (err) {
|
|
1190
|
+
if (!handled) {
|
|
1191
|
+
handled = true;
|
|
1192
|
+
guard.release(reserved);
|
|
1193
|
+
}
|
|
1194
|
+
throw err;
|
|
1195
|
+
}
|
|
935
1196
|
}
|
|
936
1197
|
controller.enqueue(chunk);
|
|
1198
|
+
},
|
|
1199
|
+
flush() {
|
|
1200
|
+
if (!handled) {
|
|
1201
|
+
handled = true;
|
|
1202
|
+
guard.release(reserved);
|
|
1203
|
+
}
|
|
1204
|
+
},
|
|
1205
|
+
cancel() {
|
|
1206
|
+
if (!handled) {
|
|
1207
|
+
handled = true;
|
|
1208
|
+
guard.release(reserved);
|
|
1209
|
+
}
|
|
937
1210
|
}
|
|
1211
|
+
// `cancel` is valid per the Streams spec and supported in Node 18+, but
|
|
1212
|
+
// TS's Transformer lib type lags and omits it — cast to keep the type check.
|
|
938
1213
|
})
|
|
939
1214
|
);
|
|
940
1215
|
return { stream: guarded, ...rest };
|