From e214a973da7d65a9e0140bf22a31ca99dd1ebf48 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Sun, 13 Oct 2024 08:39:37 +0100 Subject: [PATCH] Azure env vars can provide endpoint/token --- .../trustgraph/model/text_completion/azure/llm.py | 13 +++++++++++-- .../model/text_completion/azure_openai/llm.py | 13 +++++++++++-- .../trustgraph/model/text_completion/claude/llm.py | 4 +++- 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index ff97f644..949eeae5 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -7,6 +7,7 @@ serverless endpoint service. Input is prompt, output is response. import requests import json from prometheus_client import Histogram +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -23,6 +24,8 @@ default_subscriber = module default_temperature = 0.0 default_max_output = 4192 default_model = "AzureAI" +default_endpoint = os.getenv("AZURE_ENDPOINT") +default_token = os.getenv("AZURE_TOKEN") class Processor(ConsumerProducer): @@ -31,12 +34,18 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) - endpoint = params.get("endpoint") - token = params.get("token") + endpoint = params.get("endpoint", default_endpoint) + token = params.get("token", default_token) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) model = default_model + if endpoint is None: + raise RuntimeError("Azure endpoint not specified") + + if token is None: + raise RuntimeError("Azure token not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue, 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 d0939d1e..1dbfba27 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -8,6 +8,7 @@ import requests import json from prometheus_client import Histogram from openai import AzureOpenAI +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -24,6 +25,8 @@ default_subscriber = module default_temperature = 0.0 default_max_output = 4192 default_api = "2024-02-15-preview" +default_endpoint = os.getenv("AZURE_ENDPOINT") +default_token = os.getenv("AZURE_TOKEN") class Processor(ConsumerProducer): @@ -32,13 +35,19 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) - endpoint = params.get("endpoint") - token = params.get("token") + endpoint = params.get("endpoint", default_endpoint) + token = params.get("token", default_token) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) model = params.get("model") api = params.get("api_version", default_api) + if endpoint is None: + raise RuntimeError("Azure endpoint not specified") + + if token is None: + raise RuntimeError("Azure token not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue, diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index ad949b02..9bc37ea3 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -6,6 +6,7 @@ Input is prompt, output is response. import anthropic from prometheus_client import Histogram +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -22,6 +23,7 @@ default_subscriber = module default_model = 'claude-3-5-sonnet-20240620' default_temperature = 0.0 default_max_output = 8192 +default_api_key = os.getenv("CLAUDE_KEY") class Processor(ConsumerProducer): @@ -31,7 +33,7 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) model = params.get("model", default_model) - api_key = params.get("api_key") + api_key = params.get("api_key", default_api_key) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output)