Skip to content

Commit

Permalink
Remove the model field from the countTokens request payload (#162)
Browse files Browse the repository at this point in the history
The `model` isn't a required part of the `countTokens` request, as it's already part of the request URL.
  • Loading branch information
DellaBitta authored May 31, 2024
1 parent 83ec4ac commit 1440a05
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 1 deletion.
5 changes: 5 additions & 0 deletions .changeset/fifty-masks-leave.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@google/generative-ai": patch
---

Removed the `model` field from the internally formatted payload of `countToken` requests as it was unnecessary.
2 changes: 1 addition & 1 deletion packages/main/src/methods/count-tokens.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ export async function countTokens(
Task.COUNT_TOKENS,
apiKey,
false,
JSON.stringify({ ...params, model }),
JSON.stringify(params),
requestOptions,
);
return response.json();
Expand Down
24 changes: 24 additions & 0 deletions packages/main/test-integration/node/count-tokens.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -66,4 +66,28 @@ describe("countTokens", function () {
const response = await model.countTokens(countTokensRequest);
expect(response.totalTokens).to.equal(3);
});
it("counts tokens with GenerateContentRequest", async () => {
const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
const model = genAI.getGenerativeModel({
model: "gemini-1.5-flash-latest",
safetySettings: [
{
category: HarmCategory.HARM_CATEGORY_HARASSMENT,
threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
},
],
});
const countTokensRequest: CountTokensRequest = {
generateContentRequest: {
contents: [
{
role: "user",
parts: [{ text: "count me again with a different result" }],
},
],
},
};
const response = await model.countTokens(countTokensRequest);
expect(response.totalTokens).to.equal(8);
});
});

0 comments on commit 1440a05

Please sign in to comment.