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

Fix missing openai symbol / allow specifying base URL (#330)

Browse source
This commit is contained in:
cybermaggedon 2025-03-25 21:12:34 +00:00 • committed by GitHub
parent b1ab2938b0
commit 8efc3b281d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 11 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -24,6 +24,7 @@ default_model = 'gpt-3.5-turbo'
default_temperature = 0.0 default_temperature = 0.0
default_max_output = 4096 default_max_output = 4096
default_api_key = os.getenv("OPENAI_TOKEN") default_api_key = os.getenv("OPENAI_TOKEN")
default_base_url = os.getenv("OPENAI_BASE_URL", None)
class Processor(ConsumerProducer): class Processor(ConsumerProducer):
@ -34,6 +35,7 @@ class Processor(ConsumerProducer):
subscriber = params.get("subscriber", default_subscriber) subscriber = params.get("subscriber", default_subscriber)
model = params.get("model", default_model) model = params.get("model", default_model)
api_key = params.get("api_key", default_api_key) api_key = params.get("api_key", default_api_key)
base_url = params.get("base_url", default_base_url)
temperature = params.get("temperature", default_temperature) temperature = params.get("temperature", default_temperature)
max_output = params.get("max_output", default_max_output) max_output = params.get("max_output", default_max_output)
@ -50,6 +52,7 @@ class Processor(ConsumerProducer):
"model": model, "model": model,
"temperature": temperature, "temperature": temperature,
"max_output": max_output, "max_output": max_output,
"base_url": base_url,
} }
) )
@ -69,7 +72,7 @@ class Processor(ConsumerProducer):
self.model = model self.model = model
self.temperature = temperature self.temperature = temperature
self.max_output = max_output self.max_output = max_output
self.openai = OpenAI(api_key=api_key) self.openai = OpenAI(base_url=base_url, api_key=api_key)
print("Initialised", flush=True) print("Initialised", flush=True)
@ -132,7 +135,7 @@ class Processor(ConsumerProducer):
# FIXME: Wrong exception, don't know what this LLM throws # FIXME: Wrong exception, don't know what this LLM throws
# for a rate limit # for a rate limit
except openai.RateLimitError: except RateLimitError:
# Leave rate limit retries to the base handler # Leave rate limit retries to the base handler
raise TooManyRequests() raise TooManyRequests()
@ -180,6 +183,12 @@ class Processor(ConsumerProducer):
help=f'OpenAI API key' help=f'OpenAI API key'
) )
parser.add_argument(
'-u', '--url',
default=default_base_url,
help=f'OpenAI service base URL'
)
parser.add_argument( parser.add_argument(
'-t', '--temperature', '-t', '--temperature',
type=float, type=float,