Skip to content

Commit ed8b6da

Browse files
committed
fix: overrice countTokens in tests to use the default token counting without a safetFactor
1 parent 7963fcc commit ed8b6da

File tree

1 file changed

+31
-0
lines changed

1 file changed

+31
-0
lines changed

src/core/sliding-window/__tests__/sliding-window.spec.ts

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,37 @@ class MockApiHandler extends BaseProvider {
4343
},
4444
}
4545
}
46+
47+
// Override countTokens to use the count without factors
48+
override async countTokens(
49+
content: any[],
50+
options: {
51+
maxTokens?: number | null
52+
effectiveThreshold?: number
53+
totalTokens: number
54+
},
55+
): Promise<number> {
56+
if (content.length === 0) {
57+
return 0
58+
}
59+
60+
const { countTokens: localCountTokens } = await import("../../../utils/countTokens")
61+
62+
let totalTokens = 0
63+
for (const block of content) {
64+
if (block.type === "text") {
65+
if (block.text === "") {
66+
continue
67+
}
68+
totalTokens += await localCountTokens([block], { useWorker: true })
69+
} else if (block.type === "image") {
70+
const dataLength = block.source.data.length
71+
totalTokens += Math.ceil(Math.sqrt(dataLength)) * 1.5
72+
}
73+
}
74+
75+
return totalTokens
76+
}
4677
}
4778

4879
// Create a singleton instance for tests

0 commit comments

Comments
 (0)