diff --git a/llms/googleai/shared_test/shared_test.go b/llms/googleai/shared_test/shared_test.go index 4bdaa8bb0..2495e9810 100644 --- a/llms/googleai/shared_test/shared_test.go +++ b/llms/googleai/shared_test/shared_test.go @@ -134,6 +134,8 @@ func testMultiContentText(t *testing.T, llm llms.Model) { assert.NotEmpty(t, rsp.Choices) c1 := rsp.Choices[0] assert.Regexp(t, "(?i)dog|carnivo|canid|canine", c1.Content) + assert.Contains(t, c1.GenerationInfo, "output_tokens") + assert.NotZero(t, c1.GenerationInfo["output_tokens"]) } func testMultiContentTextUsingTextParts(t *testing.T, llm llms.Model) { @@ -317,6 +319,8 @@ func testWithStreaming(t *testing.T, llm llms.Model) { c1 := rsp.Choices[0] assert.Regexp(t, "dog|canid", strings.ToLower(c1.Content)) assert.Regexp(t, "dog|canid", strings.ToLower(sb.String())) + assert.Contains(t, c1.GenerationInfo, "output_tokens") + assert.NotZero(t, c1.GenerationInfo["output_tokens"]) } func testTools(t *testing.T, llm llms.Model) { @@ -354,6 +358,8 @@ func testTools(t *testing.T, llm llms.Model) { assert.NotEmpty(t, resp.Choices) c1 := resp.Choices[0] + assert.Contains(t, c1.GenerationInfo, "output_tokens") + assert.NotZero(t, c1.GenerationInfo["output_tokens"]) // Update chat history with assistant's response, with its tool calls. assistantResp := llms.MessageContent{ diff --git a/llms/googleai/vertex/vertex.go b/llms/googleai/vertex/vertex.go index 29e8ab107..afb50551b 100644 --- a/llms/googleai/vertex/vertex.go +++ b/llms/googleai/vertex/vertex.go @@ -113,7 +113,7 @@ func (g *Vertex) GenerateContent( } // convertCandidates converts a sequence of genai.Candidate to a response. -func convertCandidates(candidates []*genai.Candidate) (*llms.ContentResponse, error) { +func convertCandidates(candidates []*genai.Candidate, usage *genai.UsageMetadata) (*llms.ContentResponse, error) { var contentResponse llms.ContentResponse var toolCalls []llms.ToolCall @@ -149,6 +149,11 @@ func convertCandidates(candidates []*genai.Candidate) (*llms.ContentResponse, er metadata := make(map[string]any) metadata[CITATIONS] = candidate.CitationMetadata metadata[SAFETY] = candidate.SafetyRatings + if usage != nil { + metadata["input_tokens"] = usage.PromptTokenCount + metadata["output_tokens"] = usage.CandidatesTokenCount + metadata["total_tokens"] = usage.TotalTokenCount + } contentResponse.Choices = append(contentResponse.Choices, &llms.ContentChoice{ @@ -257,7 +262,7 @@ func generateFromSingleMessage( if len(resp.Candidates) == 0 { return nil, ErrNoContentInResponse } - return convertCandidates(resp.Candidates) + return convertCandidates(resp.Candidates, resp.UsageMetadata) } iter := model.GenerateContentStream(ctx, convertedParts...) return convertAndStreamFromIterator(ctx, iter, opts) @@ -296,7 +301,7 @@ func generateFromMessages( if len(resp.Candidates) == 0 { return nil, ErrNoContentInResponse } - return convertCandidates(resp.Candidates) + return convertCandidates(resp.Candidates, resp.UsageMetadata) } iter := session.SendMessageStream(ctx, reqContent.Parts...) return convertAndStreamFromIterator(ctx, iter, opts) @@ -315,6 +320,7 @@ func convertAndStreamFromIterator( candidate := &genai.Candidate{ Content: &genai.Content{}, } + var usage *genai.UsageMetadata DoStream: for { resp, err := iter.Next() @@ -339,6 +345,10 @@ DoStream: candidate.SafetyRatings = respCandidate.SafetyRatings candidate.CitationMetadata = respCandidate.CitationMetadata + if resp.UsageMetadata != nil { + usage = resp.UsageMetadata + } + for _, part := range respCandidate.Content.Parts { if text, ok := part.(genai.Text); ok { if opts.StreamingFunc(ctx, []byte(text)) != nil { @@ -348,7 +358,7 @@ DoStream: } } - return convertCandidates([]*genai.Candidate{candidate}) + return convertCandidates([]*genai.Candidate{candidate}, usage) } // convertTools converts from a list of langchaingo tools to a list of genai