From 0bcd9ffa362376cfe6af4844978b05072bc20649 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Tue, 18 Mar 2025 23:58:29 +0000 Subject: [PATCH] Fix async send issues in LLMs --- trustgraph-flow/trustgraph/model/text_completion/azure/llm.py | 4 ++-- .../trustgraph/model/text_completion/azure_openai/llm.py | 2 +- .../trustgraph/model/text_completion/cohere/llm.py | 2 +- .../trustgraph/model/text_completion/vertexai/llm.py | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index 8951b71e..33840378 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -154,7 +154,7 @@ class Processor(ConsumerProducer): print("Send response...", flush=True) r = TextCompletionResponse(response=resp, error=None, in_token=inputtokens, out_token=outputtokens, model=self.model) - await self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) except TooManyRequests: @@ -182,7 +182,7 @@ class Processor(ConsumerProducer): model=None, ) - await self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index bce44ab6..92729452 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -163,7 +163,7 @@ class Processor(ConsumerProducer): model=None, ) - await self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index 4b18f7f1..d5dab142 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -136,7 +136,7 @@ class Processor(ConsumerProducer): model=None, ) - await self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg) diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index efe63ea6..4d38c8c0 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -169,7 +169,7 @@ class Processor(ConsumerProducer): model=self.model ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) print("Done.", flush=True) @@ -202,7 +202,7 @@ class Processor(ConsumerProducer): model=None, ) - self.producer.send(r, properties={"id": id}) + await self.send(r, properties={"id": id}) self.consumer.acknowledge(msg)