apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-19 02:01:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Other env vars

Browse source
This commit is contained in:
Cyber MacGeddon 2024-10-14 14:34:25 +01:00
parent e214a973da
commit 1588d6a4f2
6 changed files with 39 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

3
trustgraph-flow/trustgraph/model/text_completion/claude/llm.py
View file

@ -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,

7
trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py
View file

@ -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,

9
trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py
View file

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

7
trustgraph-flow/trustgraph/model/text_completion/openai/llm.py
View file

@ -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,