From 577391183c41e867f4cc1400a3004d5e31b22e70 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Sun, 29 Sep 2024 11:17:08 -0700 Subject: [PATCH] Fixed VertexAI token counts --- trustgraph/model/text_completion/vertexai/llm.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph/model/text_completion/vertexai/llm.py index 41f4ff3e..c57b9fb0 100755 --- a/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph/model/text_completion/vertexai/llm.py @@ -142,14 +142,14 @@ class Processor(ConsumerProducer): with __class__.text_completion_metric.time(): - resp = self.llm.generate_content( + response = self.llm.generate_content( prompt, generation_config=self.generation_config, safety_settings=self.safety_settings ) - resp = resp.text - inputtokens = resp.usage_metadata.prompt_token_count - outputtokens = resp.usage_metadata.candidates_token_count + resp = response.text + inputtokens = int(response.usage_metadata.prompt_token_count) + outputtokens = int(response.usage_metadata.candidates_token_count) print(resp, flush=True) print(f"Input Tokens: {inputtokens}", flush=True) print(f"Output Tokens: {outputtokens}", flush=True)