From d389c4f45efad672c2aed289e289054a9ecb6b96 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Sat, 28 Sep 2024 16:45:25 -0700 Subject: [PATCH] Added support for VertexAI --- trustgraph/model/text_completion/vertexai/llm.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph/model/text_completion/vertexai/llm.py index 17ebde6c..41f4ff3e 100755 --- a/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph/model/text_completion/vertexai/llm.py @@ -122,6 +122,7 @@ class Processor(ConsumerProducer): print(f"Initialise model {model}", flush=True) self.llm = GenerativeModel(model) + self.model = model print("Initialisation complete", flush=True) @@ -147,12 +148,20 @@ class Processor(ConsumerProducer): ) resp = resp.text + inputtokens = resp.usage_metadata.prompt_token_count + outputtokens = resp.usage_metadata.candidates_token_count + print(resp, flush=True) + print(f"Input Tokens: {inputtokens}", flush=True) + print(f"Output Tokens: {outputtokens}", flush=True) print("Send response...", flush=True) r = TextCompletionResponse( error=None, response=resp, + in_token=inputtokens, + out_token=outputtokens, + model=self.model ) self.producer.send(r, properties={"id": id}) @@ -172,6 +181,9 @@ class Processor(ConsumerProducer): message = str(e), ), response=None, + in_token=None, + out_token=None, + model=None, ) self.producer.send(r, properties={"id": id}) @@ -190,6 +202,9 @@ class Processor(ConsumerProducer): message = str(e), ), response=None, + in_token=None, + out_token=None, + model=None, ) self.producer.send(r, properties={"id": id})