Other env vars

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,

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,

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,

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__)
 
-    
+    

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,

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,