Skip to content

Commit

Permalink
add tests for countTokens method in Vertex AI
Browse files Browse the repository at this point in the history
  • Loading branch information
jochenkirstaetter committed Feb 7, 2025
1 parent 72f07b9 commit e9c7ffc
Showing 1 changed file with 26 additions and 0 deletions.
26 changes: 26 additions & 0 deletions tests/Mscc.GenerativeAI/VertexAi_Gemini15Pro_Should.cs
Original file line number Diff line number Diff line change
Expand Up @@ -501,5 +501,31 @@ public async Task Generate_Content_Stream_ExpressMode()
output.WriteLine($"TotalTokenCount: {response?.UsageMetadata?.TotalTokenCount}");
}
}

[Theory]
[InlineData("How are you doing today?", 6)]
[InlineData("What kind of fish is this?", 7)]
[InlineData("Write a story about a magic backpack.", 8)]
[InlineData("Write an extended story about a magic backpack.", 9)]
public async Task Count_Tokens_Request_ExpressMode(string prompt, int expected)
{
// Arrange
var vertex = new VertexAI(apiKey: fixture.ApiKey);
var model = vertex.GenerativeModel(model: _model);
var request = new GenerateContentRequest { Contents = new List<Content>() };
request.Contents.Add(new Content
{
Role = Role.User,
Parts = new List<IPart> { new TextData { Text = prompt } }
});

// Act
var response = await model.CountTokens(request);

// Assert
response.Should().NotBeNull();
response.TotalTokens.Should().Be(expected);
output.WriteLine($"Tokens: {response?.TotalTokens}");
}
}
}

0 comments on commit e9c7ffc

Please sign in to comment.