diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index 0d050261..99cdd1c6 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -7,6 +7,7 @@ Input is prompt, output is response. Mistral is default. import boto3 import json from prometheus_client import Histogram +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -24,7 +25,8 @@ default_model = 'mistral.mistral-large-2407-v1:0' default_region = 'us-west-2' default_temperature = 0.0 default_max_output = 2048 - +default_aws_id = os.getenv("AWS_ID_KEY") +default_aws_secret = os.getenv("AWS_SECRET_KEY") class Processor(ConsumerProducer): @@ -34,12 +36,18 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) model = params.get("model", default_model) - aws_id = params.get("aws_id_key") - aws_secret = params.get("aws_secret") + aws_id = params.get("aws_id_key", default_aws_id) + aws_secret = params.get("aws_secret", default_aws_secret) aws_region = params.get("aws_region", default_region) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) + if aws_id is None: + raise RuntimeError("AWS ID not specified") + + if aws_secret is None: + raise RuntimeError("AWS secret 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 9bc37ea3..08e2828d 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -37,6 +37,9 @@ class Processor(ConsumerProducer): temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) + if api_key is None: + raise RuntimeError("Claude API key not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue, diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index 4c64e8b6..8e1f4f7c 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -6,6 +6,7 @@ Input is prompt, output is response. import cohere from prometheus_client import Histogram +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -21,6 +22,7 @@ default_output_queue = text_completion_response_queue default_subscriber = module default_model = 'c4ai-aya-23-8b' default_temperature = 0.0 +default_api_key = os.getenv("COHERE_KEY") class Processor(ConsumerProducer): @@ -30,9 +32,12 @@ 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) + if api_key is None: + raise RuntimeError("Cohere API key not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue, diff --git a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py index baa00b04..47a75927 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py @@ -7,6 +7,7 @@ Input is prompt, output is response. import google.generativeai as genai from google.generativeai.types import HarmCategory, HarmBlockThreshold from prometheus_client import Histogram +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -23,6 +24,7 @@ default_subscriber = module default_model = 'gemini-1.5-flash-002' default_temperature = 0.0 default_max_output = 8192 +default_api_key = os.getenv("GOOGLE_AI_STUDIO_KEY") class Processor(ConsumerProducer): @@ -32,10 +34,13 @@ 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) + if api_key is None: + raise RuntimeError("Google AI Studio API key not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue, @@ -211,4 +216,4 @@ def run(): Processor.start(module, __doc__) - \ No newline at end of file + diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index 5d259e7e..b63ad43b 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -6,6 +6,7 @@ Input is prompt, output is response. from openai import OpenAI 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 = 'gpt-3.5-turbo' default_temperature = 0.0 default_max_output = 4096 +default_api_key = os.getenv("OPENAI_KEY") class Processor(ConsumerProducer): @@ -31,10 +33,13 @@ 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) + if api_key is None: + raise RuntimeError("OpenAI API key not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue, diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index c57b9fb0..76d517a1 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -7,6 +7,7 @@ Google Cloud. Input is prompt, output is response. import vertexai import time from prometheus_client import Histogram +import os from google.oauth2 import service_account import google @@ -38,6 +39,7 @@ default_model = 'gemini-1.0-pro-001' default_region = 'us-central1' default_temperature = 0.0 default_max_output = 8192 +default_private_key = os.getenv("VERTEXAI_KEY") class Processor(ConsumerProducer): @@ -48,10 +50,13 @@ class Processor(ConsumerProducer): subscriber = params.get("subscriber", default_subscriber) region = params.get("region", default_region) model = params.get("model", default_model) - private_key = params.get("private_key") + private_key = params.get("private_key", default_private_key) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) + if private_key is None: + raise RuntimeError("Private key file not specified") + super(Processor, self).__init__( **params | { "input_queue": input_queue,