Add support to set up openai organizations (#2514)

Add support for defining the organization of OpenAI, similarly to what
is done in the reference code below:

```
import os
import openai
openai.organization = os.getenv("OPENAI_ORGANIZATION")
openai.api_key = os.getenv("OPENAI_API_KEY")
```
This commit is contained in:
Tiago De Gaspari 2023-04-07 02:23:16 -03:00 committed by GitHub
parent 629fda3957
commit c2f21a519f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 32 additions and 0 deletions

View File

@ -38,10 +38,15 @@ class OpenAIModerationChain(Chain):
openai_api_key = get_from_dict_or_env( openai_api_key = get_from_dict_or_env(
values, "openai_api_key", "OPENAI_API_KEY" values, "openai_api_key", "OPENAI_API_KEY"
) )
openai_organization = get_from_dict_or_env(
values, "openai_organization", "OPENAI_ORGANIZATION", default=None
)
try: try:
import openai import openai
openai.api_key = openai_api_key openai.api_key = openai_api_key
if openai_organization:
openai.organization = openai_organization
values["client"] = openai.Moderation values["client"] = openai.Moderation
except ImportError: except ImportError:
raise ValueError( raise ValueError(

View File

@ -68,6 +68,9 @@ class AzureChatOpenAI(ChatOpenAI):
"openai_api_type", "openai_api_type",
"OPENAI_API_TYPE", "OPENAI_API_TYPE",
) )
openai_organization = get_from_dict_or_env(
values, "openai_organization", "OPENAI_ORGANIZATION", default=None
)
try: try:
import openai import openai
@ -75,6 +78,8 @@ class AzureChatOpenAI(ChatOpenAI):
openai.api_base = openai_api_base openai.api_base = openai_api_base
openai.api_version = openai_api_version openai.api_version = openai_api_version
openai.api_key = openai_api_key openai.api_key = openai_api_key
if openai_organization:
openai.organization = openai_organization
except ImportError: except ImportError:
raise ValueError( raise ValueError(
"Could not import openai python package. " "Could not import openai python package. "

View File

@ -151,10 +151,15 @@ class ChatOpenAI(BaseChatModel):
openai_api_key = get_from_dict_or_env( openai_api_key = get_from_dict_or_env(
values, "openai_api_key", "OPENAI_API_KEY" values, "openai_api_key", "OPENAI_API_KEY"
) )
openai_organization = get_from_dict_or_env(
values, "openai_organization", "OPENAI_ORGANIZATION", default=None
)
try: try:
import openai import openai
openai.api_key = openai_api_key openai.api_key = openai_api_key
if openai_organization:
openai.organization = openai_organization
except ImportError: except ImportError:
raise ValueError( raise ValueError(
"Could not import openai python package. " "Could not import openai python package. "

View File

@ -159,10 +159,15 @@ class OpenAIEmbeddings(BaseModel, Embeddings):
openai_api_key = get_from_dict_or_env( openai_api_key = get_from_dict_or_env(
values, "openai_api_key", "OPENAI_API_KEY" values, "openai_api_key", "OPENAI_API_KEY"
) )
openai_organization = get_from_dict_or_env(
values, "openai_organization", "OPENAI_ORGANIZATION", default=None
)
try: try:
import openai import openai
openai.api_key = openai_api_key openai.api_key = openai_api_key
if openai_organization:
openai.organization = openai_organization
values["client"] = openai.Embedding values["client"] = openai.Embedding
except ImportError: except ImportError:
raise ValueError( raise ValueError(

View File

@ -204,10 +204,17 @@ class BaseOpenAI(BaseLLM):
openai_api_key = get_from_dict_or_env( openai_api_key = get_from_dict_or_env(
values, "openai_api_key", "OPENAI_API_KEY" values, "openai_api_key", "OPENAI_API_KEY"
) )
openai_organization = get_from_dict_or_env(
values, "openai_organization", "OPENAI_ORGANIZATION", default=None
)
try: try:
import openai import openai
openai.api_key = openai_api_key openai.api_key = openai_api_key
if openai_organization:
print("USING ORGANIZATION: ")
print(openai_organization)
openai.organization = openai_organization
values["client"] = openai.Completion values["client"] = openai.Completion
except ImportError: except ImportError:
raise ValueError( raise ValueError(
@ -588,10 +595,15 @@ class OpenAIChat(BaseLLM):
openai_api_key = get_from_dict_or_env( openai_api_key = get_from_dict_or_env(
values, "openai_api_key", "OPENAI_API_KEY" values, "openai_api_key", "OPENAI_API_KEY"
) )
openai_organization = get_from_dict_or_env(
values, "openai_organization", "OPENAI_ORGANIZATION", default=None
)
try: try:
import openai import openai
openai.api_key = openai_api_key openai.api_key = openai_api_key
if openai_organization:
openai.organization = openai_organization
except ImportError: except ImportError:
raise ValueError( raise ValueError(
"Could not import openai python package. " "Could not import openai python package. "