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

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)
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})