diff --git a/python/src/typechat/_internal/model.py b/python/src/typechat/_internal/model.py index 471404f0..42e412a1 100644 --- a/python/src/typechat/_internal/model.py +++ b/python/src/typechat/_internal/model.py @@ -180,7 +180,7 @@ def required_var(name: str) -> str: api_key = required_var("OPENAI_API_KEY") model = required_var("OPENAI_MODEL") endpoint = vals.get("OPENAI_ENDPOINT", None) or "https://api.openai.com/v1/chat/completions" - org = vals.get("OPENAI_ORG", None) or "" + org = vals.get("OPENAI_ORGANIZATION", None) or vals.get("OPENAI_ORG", None) or "" return create_openai_language_model(api_key, model, endpoint, org) elif "AZURE_OPENAI_API_KEY" in vals: diff --git a/python/tests/test_create_language_model.py b/python/tests/test_create_language_model.py new file mode 100644 index 00000000..62dbe01c --- /dev/null +++ b/python/tests/test_create_language_model.py @@ -0,0 +1,40 @@ +""" +Tests for create_language_model's handling of the OpenAI organization variables. + +These only inspect the constructed model, so no request is sent. +""" + +import typechat + +def test_openai_organization_is_sent_as_organization_header(): + model = typechat.create_language_model( + { + "OPENAI_MODEL": "gpt-4o", + "OPENAI_API_KEY": "openai-secret", + "OPENAI_ORGANIZATION": "org-123", + } + ) + assert model.headers["OpenAI-Organization"] == "org-123" + + +def test_openai_org_is_still_sent_as_organization_header(): + model = typechat.create_language_model( + { + "OPENAI_MODEL": "gpt-4o", + "OPENAI_API_KEY": "openai-secret", + "OPENAI_ORG": "org-legacy", + } + ) + assert model.headers["OpenAI-Organization"] == "org-legacy" + + +def test_openai_organization_takes_precedence_over_openai_org(): + model = typechat.create_language_model( + { + "OPENAI_MODEL": "gpt-4o", + "OPENAI_API_KEY": "openai-secret", + "OPENAI_ORGANIZATION": "org-123", + "OPENAI_ORG": "org-legacy", + } + ) + assert model.headers["OpenAI-Organization"] == "org-123"