diff --git a/semantic_router/encoders/openai.py b/semantic_router/encoders/openai.py index 7957f913b342b8fc8a13284783f38d3ca62ce176..14fda1c99adee1a60a9f03621ade6f93151410db 100644 --- a/semantic_router/encoders/openai.py +++ b/semantic_router/encoders/openai.py @@ -20,6 +20,7 @@ class OpenAIEncoder(BaseEncoder): self, name: Optional[str] = None, openai_api_key: Optional[str] = None, + openai_org_id: Optional[str] = None, score_threshold: float = 0.82, dimensions: Union[int, NotGiven] = NotGiven(), ): @@ -27,10 +28,11 @@ class OpenAIEncoder(BaseEncoder): name = os.getenv("OPENAI_MODEL_NAME", "text-embedding-ada-002") super().__init__(name=name, score_threshold=score_threshold) api_key = openai_api_key or os.getenv("OPENAI_API_KEY") + openai_org_id = openai_org_id or os.getenv("OPENAI_ORG_ID") if api_key is None: raise ValueError("OpenAI API key cannot be 'None'.") try: - self.client = openai.Client(api_key=api_key) + self.client = openai.Client(api_key=api_key, organization=openai_org_id) except Exception as e: raise ValueError( f"OpenAI API client failed to initialize. Error: {e}" diff --git a/tests/unit/encoders/test_openai.py b/tests/unit/encoders/test_openai.py index 4679ee939f7d4b494150a8a52f4b4c33a0e6c8db..508e9e9e197a8a0b87fb179554d43c18686c1b44 100644 --- a/tests/unit/encoders/test_openai.py +++ b/tests/unit/encoders/test_openai.py @@ -14,7 +14,9 @@ def openai_encoder(mocker): class TestOpenAIEncoder: def test_openai_encoder_init_success(self, mocker): - mocker.patch("os.getenv", return_value="fake-api-key") + # -- Mock the return value of os.getenv 3 times: model name, api key and org ID + side_effect = ["fake-model-name", "fake-api-key", "fake-org-id"] + mocker.patch("os.getenv", side_effect=side_effect) encoder = OpenAIEncoder() assert encoder.client is not None