apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-18 17:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added support for VertexAI

Browse source
This commit is contained in:
JackColquitt 2024-09-28 16:45:25 -07:00
parent 57c506b32e
commit d389c4f45e
1 changed files with 15 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

15
trustgraph/model/text_completion/vertexai/llm.py
View file

@ -122,6 +122,7 @@ class Processor(ConsumerProducer):
print(f"Initialise model {model}", flush=True) print(f"Initialise model {model}", flush=True)
self.llm = GenerativeModel(model) self.llm = GenerativeModel(model)
self.model = model
print("Initialisation complete", flush=True) print("Initialisation complete", flush=True)
@ -147,12 +148,20 @@ class Processor(ConsumerProducer):
) )
resp = resp.text 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) print("Send response...", flush=True)
r = TextCompletionResponse( r = TextCompletionResponse(
error=None, error=None,
response=resp, response=resp,
in_token=inputtokens,
out_token=outputtokens,
model=self.model
) )
self.producer.send(r, properties={"id": id}) self.producer.send(r, properties={"id": id})
@ -172,6 +181,9 @@ class Processor(ConsumerProducer):
message = str(e), message = str(e),
), ),
response=None, response=None,
in_token=None,
out_token=None,
model=None,
) )
self.producer.send(r, properties={"id": id}) self.producer.send(r, properties={"id": id})
@ -190,6 +202,9 @@ class Processor(ConsumerProducer):
message = str(e), message = str(e),
), ),
response=None, response=None,
in_token=None,
out_token=None,
model=None,
) )
self.producer.send(r, properties={"id": id}) self.producer.send(r, properties={"id": id})