diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index ebfae9ed..590c2e3f 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -24,6 +24,7 @@ default_model = 'gpt-3.5-turbo' default_temperature = 0.0 default_max_output = 4096 default_api_key = os.getenv("OPENAI_TOKEN") +default_base_url = os.getenv("OPENAI_BASE_URL", None) class Processor(ConsumerProducer): @@ -34,6 +35,7 @@ class Processor(ConsumerProducer): subscriber = params.get("subscriber", default_subscriber) model = params.get("model", default_model) api_key = params.get("api_key", default_api_key) + base_url = params.get("base_url", default_base_url) temperature = params.get("temperature", default_temperature) max_output = params.get("max_output", default_max_output) @@ -50,6 +52,7 @@ class Processor(ConsumerProducer): "model": model, "temperature": temperature, "max_output": max_output, + "base_url": base_url, } ) @@ -69,7 +72,7 @@ class Processor(ConsumerProducer): self.model = model self.temperature = temperature 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) @@ -132,7 +135,7 @@ class Processor(ConsumerProducer): # FIXME: Wrong exception, don't know what this LLM throws # for a rate limit - except openai.RateLimitError: + except RateLimitError: # Leave rate limit retries to the base handler raise TooManyRequests() @@ -180,6 +183,12 @@ class Processor(ConsumerProducer): help=f'OpenAI API key' ) + parser.add_argument( + '-u', '--url', + default=default_base_url, + help=f'OpenAI service base URL' + ) + parser.add_argument( '-t', '--temperature', type=float,