From 294c2d157473fa035d6d930d40053fd291587931 Mon Sep 17 00:00:00 2001 From: bobo Date: Thu, 25 Jan 2024 23:33:34 +0800 Subject: [PATCH] Add environment variables for OpenAI API key and base URL --- unit_tests/openai_util.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/unit_tests/openai_util.py b/unit_tests/openai_util.py index ea099f9..14f5b73 100644 --- a/unit_tests/openai_util.py +++ b/unit_tests/openai_util.py @@ -1,3 +1,4 @@ +import os from typing import Optional from openai import OpenAI, Stream @@ -13,7 +14,10 @@ def create_chat_completion_chunks( """ Create streaming responses. """ - _client = client or OpenAI() + _client = client or OpenAI( + api_key=os.environ.get("OPENAI_API_KEY", None), + base_url=os.environ.get("OPENAI_API_BASE", None), + ) # Force to use streaming kwargs["stream"] = True @@ -36,7 +40,10 @@ def create_chat_completion_content(client: Optional[OpenAI] = None, **kwargs) -> This is a replacement of creating non-streaming chat completion. """ - _client = client or OpenAI() + _client = client or OpenAI( + api_key=os.environ.get("OPENAI_API_KEY", None), + base_url=os.environ.get("OPENAI_API_BASE", None), + ) # Force to use streaming kwargs["stream"] = True