From c2f21a519f02ca06c8cf8db9e293ad7b5a96ec09 Mon Sep 17 00:00:00 2001
From: Tiago De Gaspari <tiago.gaspari@ae.studio>
Date: Fri, 7 Apr 2023 02:23:16 -0300
Subject: [PATCH] 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")
```
---
 langchain/chains/moderation.py        |  5 +++++
 langchain/chat_models/azure_openai.py |  5 +++++
 langchain/chat_models/openai.py       |  5 +++++
 langchain/embeddings/openai.py        |  5 +++++
 langchain/llms/openai.py              | 12 ++++++++++++
 5 files changed, 32 insertions(+)

diff --git a/langchain/chains/moderation.py b/langchain/chains/moderation.py
index de02ee0b1d..f0ec4afd80 100644
--- a/langchain/chains/moderation.py
+++ b/langchain/chains/moderation.py
@@ -38,10 +38,15 @@ class OpenAIModerationChain(Chain):
         openai_api_key = get_from_dict_or_env(
             values, "openai_api_key", "OPENAI_API_KEY"
         )
+        openai_organization = get_from_dict_or_env(
+            values, "openai_organization", "OPENAI_ORGANIZATION", default=None
+        )
         try:
             import openai
 
             openai.api_key = openai_api_key
+            if openai_organization:
+                openai.organization = openai_organization
             values["client"] = openai.Moderation
         except ImportError:
             raise ValueError(
diff --git a/langchain/chat_models/azure_openai.py b/langchain/chat_models/azure_openai.py
index f4765db09f..849a82ac70 100644
--- a/langchain/chat_models/azure_openai.py
+++ b/langchain/chat_models/azure_openai.py
@@ -68,6 +68,9 @@ class AzureChatOpenAI(ChatOpenAI):
             "openai_api_type",
             "OPENAI_API_TYPE",
         )
+        openai_organization = get_from_dict_or_env(
+            values, "openai_organization", "OPENAI_ORGANIZATION", default=None
+        )
         try:
             import openai
 
@@ -75,6 +78,8 @@ class AzureChatOpenAI(ChatOpenAI):
             openai.api_base = openai_api_base
             openai.api_version = openai_api_version
             openai.api_key = openai_api_key
+            if openai_organization:
+                openai.organization = openai_organization
         except ImportError:
             raise ValueError(
                 "Could not import openai python package. "
diff --git a/langchain/chat_models/openai.py b/langchain/chat_models/openai.py
index 2d1a0282aa..f05c9ec46c 100644
--- a/langchain/chat_models/openai.py
+++ b/langchain/chat_models/openai.py
@@ -151,10 +151,15 @@ class ChatOpenAI(BaseChatModel):
         openai_api_key = get_from_dict_or_env(
             values, "openai_api_key", "OPENAI_API_KEY"
         )
+        openai_organization = get_from_dict_or_env(
+            values, "openai_organization", "OPENAI_ORGANIZATION", default=None
+        )
         try:
             import openai
 
             openai.api_key = openai_api_key
+            if openai_organization:
+                openai.organization = openai_organization
         except ImportError:
             raise ValueError(
                 "Could not import openai python package. "
diff --git a/langchain/embeddings/openai.py b/langchain/embeddings/openai.py
index 20351fd439..fd57d281ea 100644
--- a/langchain/embeddings/openai.py
+++ b/langchain/embeddings/openai.py
@@ -159,10 +159,15 @@ class OpenAIEmbeddings(BaseModel, Embeddings):
         openai_api_key = get_from_dict_or_env(
             values, "openai_api_key", "OPENAI_API_KEY"
         )
+        openai_organization = get_from_dict_or_env(
+            values, "openai_organization", "OPENAI_ORGANIZATION", default=None
+        )
         try:
             import openai
 
             openai.api_key = openai_api_key
+            if openai_organization:
+                openai.organization = openai_organization
             values["client"] = openai.Embedding
         except ImportError:
             raise ValueError(
diff --git a/langchain/llms/openai.py b/langchain/llms/openai.py
index d4d24bda18..66c2d70a58 100644
--- a/langchain/llms/openai.py
+++ b/langchain/llms/openai.py
@@ -204,10 +204,17 @@ class BaseOpenAI(BaseLLM):
         openai_api_key = get_from_dict_or_env(
             values, "openai_api_key", "OPENAI_API_KEY"
         )
+        openai_organization = get_from_dict_or_env(
+            values, "openai_organization", "OPENAI_ORGANIZATION", default=None
+        )
         try:
             import openai
 
             openai.api_key = openai_api_key
+            if openai_organization:
+                print("USING ORGANIZATION: ")
+                print(openai_organization)
+                openai.organization = openai_organization
             values["client"] = openai.Completion
         except ImportError:
             raise ValueError(
@@ -588,10 +595,15 @@ class OpenAIChat(BaseLLM):
         openai_api_key = get_from_dict_or_env(
             values, "openai_api_key", "OPENAI_API_KEY"
         )
+        openai_organization = get_from_dict_or_env(
+            values, "openai_organization", "OPENAI_ORGANIZATION", default=None
+        )
         try:
             import openai
 
             openai.api_key = openai_api_key
+            if openai_organization:
+                openai.organization = openai_organization
         except ImportError:
             raise ValueError(
                 "Could not import openai python package. "