Skip to content

Commit 91da387

Browse files
jairad26jairad26
authored
[CHORE] Add support for non-prefixed ef env vars in python (#5764)
## Description of changes _Summarize the changes made by this PR._ - Improvements & Bug fixes - This PR allows users to set env vars without the CHROMA_ prefix on embedding functions in python, while still supporting the default env var as fallback - New functionality - ... <img width="490" height="421" alt="image" src="https://github.com/user-attachments/assets/2edc560e-e8f4-41b3-9a64-612488d8a3e1" /> <img width="916" height="160" alt="image" src="https://github.com/user-attachments/assets/45223204-1cf1-4fc0-9f20-3b5241744b61" /> <img width="727" height="161" alt="image" src="https://github.com/user-attachments/assets/b78f3f4f-5e6d-4ae5-95d8-90c009235722" /> <img width="619" height="195" alt="image" src="https://github.com/user-attachments/assets/2c32f124-6a3d-48b1-b771-c5441f47c6b9" /> <img width="829" height="250" alt="image" src="https://github.com/user-attachments/assets/7ccb1f98-6877-4f7b-8457-5b40337c0d97" /> for cloudflare, jina -> used the same naming convention as JS client ## Test plan _How are these changes tested?_ - [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust ## Migration plan _Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_ ## Observability plan _What is the plan to instrument and monitor this change?_ ## Documentation Changes _Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_
1 parent 36af0d1 commit 91da387

10 files changed

+107
-38
lines changed

‎chromadb/utils/embedding_functions/baseten_embedding_function.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,16 @@ def __init__(
3535
DeprecationWarning,
3636
)
3737

38-
self.api_key_env_var = api_key_env_var
38+
if os.getenv("BASETEN_API_KEY") is not None:
39+
self.api_key_env_var = "BASETEN_API_KEY"
40+
else:
41+
self.api_key_env_var = api_key_env_var
42+
3943
# Prioritize api_key argument, then environment variable
40-
resolved_api_key = api_key or os.getenv(api_key_env_var)
44+
resolved_api_key = api_key or os.getenv(self.api_key_env_var)
4145
if not resolved_api_key:
4246
raise ValueError(
43-
f"API key not provided and {api_key_env_var} environment variable is not set."
47+
f"API key not provided and {self.api_key_env_var} environment variable is not set."
4448
)
4549
self.api_key = resolved_api_key
4650
if not api_base:

‎chromadb/utils/embedding_functions/cloudflare_workers_ai_embedding_function.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,19 @@ def __init__(
5353
)
5454
self.model_name = model_name
5555
self.account_id = account_id
56-
self.api_key_env_var = api_key_env_var
57-
self.api_key = api_key or os.getenv(api_key_env_var)
56+
57+
if os.getenv("CLOUDFLARE_API_KEY") is not None:
58+
self.api_key_env_var = "CLOUDFLARE_API_KEY"
59+
else:
60+
self.api_key_env_var = api_key_env_var
61+
62+
self.api_key = api_key or os.getenv(self.api_key_env_var)
5863
self.gateway_id = gateway_id
5964

6065
if not self.api_key:
61-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
66+
raise ValueError(
67+
f"The {self.api_key_env_var} environment variable is not set."
68+
)
6269

6370
if self.gateway_id:
6471
self._api_url = f"{GATEWAY_BASE_URL}/{self.account_id}/{self.gateway_id}/workers-ai/{self.model_name}"

‎chromadb/utils/embedding_functions/cohere_embedding_function.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,10 +43,16 @@ def __init__(
4343
"Please use environment variables via api_key_env_var for persistent storage.",
4444
DeprecationWarning,
4545
)
46-
self.api_key_env_var = api_key_env_var
47-
self.api_key = api_key or os.getenv(api_key_env_var)
46+
if os.getenv("COHERE_API_KEY") is not None:
47+
self.api_key_env_var = "COHERE_API_KEY"
48+
else:
49+
self.api_key_env_var = api_key_env_var
50+
51+
self.api_key = api_key or os.getenv(self.api_key_env_var)
4852
if not self.api_key:
49-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
53+
raise ValueError(
54+
f"The {self.api_key_env_var} environment variable is not set."
55+
)
5056

5157
self.model_name = model_name
5258

‎chromadb/utils/embedding_functions/google_embedding_function.py‎

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,16 @@ def __init__(
3838
"Please use environment variables via api_key_env_var for persistent storage.",
3939
DeprecationWarning,
4040
)
41-
self.api_key_env_var = api_key_env_var
42-
self.api_key = api_key or os.getenv(api_key_env_var)
41+
if os.getenv("GOOGLE_API_KEY") is not None:
42+
self.api_key_env_var = "GOOGLE_API_KEY"
43+
else:
44+
self.api_key_env_var = api_key_env_var
45+
46+
self.api_key = api_key or os.getenv(self.api_key_env_var)
4347
if not self.api_key:
44-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
48+
raise ValueError(
49+
f"The {self.api_key_env_var} environment variable is not set."
50+
)
4551

4652
self.model_name = model_name
4753

@@ -154,10 +160,16 @@ def __init__(
154160
"Please use environment variables via api_key_env_var for persistent storage.",
155161
DeprecationWarning,
156162
)
157-
self.api_key_env_var = api_key_env_var
158-
self.api_key = api_key or os.getenv(api_key_env_var)
163+
if os.getenv("GOOGLE_API_KEY") is not None:
164+
self.api_key_env_var = "GOOGLE_API_KEY"
165+
else:
166+
self.api_key_env_var = api_key_env_var
167+
168+
self.api_key = api_key or os.getenv(self.api_key_env_var)
159169
if not self.api_key:
160-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
170+
raise ValueError(
171+
f"The {self.api_key_env_var} environment variable is not set."
172+
)
161173

162174
self.model_name = model_name
163175
self.task_type = task_type
@@ -289,10 +301,16 @@ def __init__(
289301
"Please use environment variables via api_key_env_var for persistent storage.",
290302
DeprecationWarning,
291303
)
292-
self.api_key_env_var = api_key_env_var
293-
self.api_key = api_key or os.getenv(api_key_env_var)
304+
if os.getenv("GOOGLE_API_KEY") is not None:
305+
self.api_key_env_var = "GOOGLE_API_KEY"
306+
else:
307+
self.api_key_env_var = api_key_env_var
308+
309+
self.api_key = api_key or os.getenv(self.api_key_env_var)
294310
if not self.api_key:
295-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
311+
raise ValueError(
312+
f"The {self.api_key_env_var} environment variable is not set."
313+
)
296314

297315
self.model_name = model_name
298316
self.project_id = project_id

‎chromadb/utils/embedding_functions/huggingface_embedding_function.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,10 +40,16 @@ def __init__(
4040
"Please use environment variables via api_key_env_var for persistent storage.",
4141
DeprecationWarning,
4242
)
43-
self.api_key_env_var = api_key_env_var
44-
self.api_key = api_key or os.getenv(api_key_env_var)
43+
if os.getenv("HUGGINGFACE_API_KEY") is not None:
44+
self.api_key_env_var = "HUGGINGFACE_API_KEY"
45+
else:
46+
self.api_key_env_var = api_key_env_var
47+
48+
self.api_key = api_key or os.getenv(self.api_key_env_var)
4549
if not self.api_key:
46-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
50+
raise ValueError(
51+
f"The {self.api_key_env_var} environment variable is not set."
52+
)
4753

4854
self.model_name = model_name
4955

@@ -160,6 +166,9 @@ def __init__(
160166
self.url = url
161167

162168
self.api_key_env_var = api_key_env_var
169+
if os.getenv("HUGGINGFACE_API_KEY") is not None:
170+
self.api_key_env_var = "HUGGINGFACE_API_KEY"
171+
163172
if self.api_key_env_var is not None:
164173
self.api_key = api_key or os.getenv(self.api_key_env_var)
165174
else:

‎chromadb/utils/embedding_functions/jina_embedding_function.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,10 +81,16 @@ def __init__(
8181
DeprecationWarning,
8282
)
8383

84-
self.api_key_env_var = api_key_env_var
85-
self.api_key = api_key or os.getenv(api_key_env_var)
84+
if os.getenv("JINA_API_KEY") is not None:
85+
self.api_key_env_var = "JINA_API_KEY"
86+
else:
87+
self.api_key_env_var = api_key_env_var
88+
89+
self.api_key = api_key or os.getenv(self.api_key_env_var)
8690
if not self.api_key:
87-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
91+
raise ValueError(
92+
f"The {self.api_key_env_var} environment variable is not set."
93+
)
8894

8995
self.model_name = model_name
9096

‎chromadb/utils/embedding_functions/openai_embedding_function.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,10 +57,16 @@ def __init__(
5757
DeprecationWarning,
5858
)
5959

60-
self.api_key_env_var = api_key_env_var
61-
self.api_key = api_key or os.getenv(api_key_env_var)
60+
if os.getenv("OPENAI_API_KEY") is not None:
61+
self.api_key_env_var = "OPENAI_API_KEY"
62+
else:
63+
self.api_key_env_var = api_key_env_var
64+
65+
self.api_key = api_key or os.getenv(self.api_key_env_var)
6266
if not self.api_key:
63-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
67+
raise ValueError(
68+
f"The {self.api_key_env_var} environment variable is not set."
69+
)
6470

6571
self.model_name = model_name
6672
self.organization_id = organization_id

‎chromadb/utils/embedding_functions/roboflow_embedding_function.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,10 +45,16 @@ def __init__(
4545
"Please use environment variables via api_key_env_var for persistent storage.",
4646
DeprecationWarning,
4747
)
48-
self.api_key_env_var = api_key_env_var
49-
self.api_key = api_key or os.getenv(api_key_env_var)
48+
if os.getenv("ROBOFLOW_API_KEY") is not None:
49+
self.api_key_env_var = "ROBOFLOW_API_KEY"
50+
else:
51+
self.api_key_env_var = api_key_env_var
52+
53+
self.api_key = api_key or os.getenv(self.api_key_env_var)
5054
if not self.api_key:
51-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
55+
raise ValueError(
56+
f"The {self.api_key_env_var} environment variable is not set."
57+
)
5258

5359
self.api_url = api_url
5460

‎chromadb/utils/embedding_functions/together_ai_embedding_function.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -48,15 +48,16 @@ def __init__(
4848
)
4949

5050
self.model_name = model_name
51-
self.api_key = api_key
52-
self.api_key_env_var = api_key_env_var
5351

54-
if not self.api_key:
55-
self.api_key = os.getenv(self.api_key_env_var)
52+
if os.getenv("TOGETHER_API_KEY") is not None:
53+
self.api_key_env_var = "TOGETHER_API_KEY"
54+
else:
55+
self.api_key_env_var = api_key_env_var
5656

57+
self.api_key = api_key or os.getenv(self.api_key_env_var)
5758
if not self.api_key:
5859
raise ValueError(
59-
f"API key not found in environment variable {self.api_key_env_var}"
60+
f"The {self.api_key_env_var} environment variable is not set."
6061
)
6162

6263
self._session = httpx.Client()

‎chromadb/utils/embedding_functions/voyageai_embedding_function.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,16 @@ def __init__(
4747
DeprecationWarning,
4848
)
4949

50-
self.api_key_env_var = api_key_env_var
51-
self.api_key = api_key or os.getenv(api_key_env_var)
50+
if os.getenv("VOYAGE_API_KEY") is not None:
51+
self.api_key_env_var = "VOYAGE_API_KEY"
52+
else:
53+
self.api_key_env_var = api_key_env_var
54+
55+
self.api_key = api_key or os.getenv(self.api_key_env_var)
5256
if not self.api_key:
53-
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
57+
raise ValueError(
58+
f"The {self.api_key_env_var} environment variable is not set."
59+
)
5460

5561
self.model_name = model_name
5662
self.input_type = input_type

0 commit comments

Comments
 (0)