From c23ca68732931c99da9203f1789605e7ff8974d0 Mon Sep 17 00:00:00 2001 From: JackColquitt Date: Sun, 29 Sep 2024 12:34:01 -0700 Subject: [PATCH] Fix for prometheus for AzureAI --- trustgraph/model/text_completion/azure/llm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/trustgraph/model/text_completion/azure/llm.py b/trustgraph/model/text_completion/azure/llm.py index 72c83e6c..ff97f644 100755 --- a/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph/model/text_completion/azure/llm.py @@ -22,6 +22,7 @@ default_output_queue = text_completion_response_queue default_subscriber = module default_temperature = 0.0 default_max_output = 4192 +default_model = "AzureAI" class Processor(ConsumerProducer): @@ -34,6 +35,7 @@ class Processor(ConsumerProducer): token = params.get("token") temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) + model = default_model super(Processor, self).__init__( **params | { @@ -44,6 +46,7 @@ class Processor(ConsumerProducer): "output_schema": TextCompletionResponse, "temperature": temperature, "max_output": max_output, + "model": model, } ) @@ -64,7 +67,7 @@ class Processor(ConsumerProducer): self.token = token self.temperature = temperature self.max_output = max_output - self.model = "AzureAI" + self.model = model def build_prompt(self, system, content):