fix: don't remove all API keys in ElysiaKeyManager

docs/creating_tools.md

@@ -136,3 +136,33 @@ All optional arguments you can pass to the `@tool` decorator are:
 - `branch_id` (`str`): the ID of the branch on the tree to add the tool to.
 - `status` (`str`): a custom message to display whilst the tool is running.
 - `end` (`bool`): when `True`, this tool can be the end of the conversation if the decision agent decides it should end after the completion of this tool.
+
+## Environment Variables
+
+Environment variables (stored in `.env` and accessed via `os.getenv(...)` or `os.environ`) are passed down to tools - with some exceptions. Any environment variables whose keys end with the following strings:
+
+- `"api_key"`
+- `"apikey"`
+- `"api_version"`
+- `"api_base"`
+- `"_account_id"`
+- `"_jwt"`
+- `"_access_key_id"`
+- `"_secret_access_key"`
+- `"_region_name"`
+- `"_token"`
+
+Will **NOT** be passed down to tools. This is to ensure that API keys are differentiated between the `Settings` object and the actual `.env` file. For example, if these *were* included in all tool calls, then any LLM calls would use what is in the `.env`, instead of what has been configured in the tree. This is due to the way LiteLLM processes LLM requests via API keys in the environment.
+
+**If you need environment variables like this, or notice your environment variables are going missing**, then consider renaming them (e.g. `my_api_key` -> `my_api_key_`) or including them in the `api_keys` variable in the `Settings` object, via
+
+```python
+from elysia import settings
+settings = Settings()
+settings.configure(
+    api_keys = {
+        "MY_API_KEY": "..."
+    }
+)
+```
+The `MY_API_KEY` will then accessible from `os.environ`. Or, in the Elysia web app, within your API keys section of the config.

elysia/api/utils/config.py

@@ -122,14 +122,15 @@ def __init__(self, logger: Logger):
         }
         self.save_location_wcd_url: str = os.getenv("WCD_URL", "")
         self.save_location_wcd_api_key: str = os.getenv("WCD_API_KEY", "")
-        self.save_location_weaviate_is_local: bool = os.getenv(
-            "WEAVIATE_IS_LOCAL", "False"
-        ) == "True"
+        self.save_location_weaviate_is_local: bool = (
+            os.getenv("WEAVIATE_IS_LOCAL", "False") == "True"
+        )
         self.save_location_local_weaviate_port: int = int(
             os.getenv("LOCAL_WEAVIATE_PORT", 8080)
         )
-        self.save_location_local_weaviate_grpc_port: int = int(os.getenv("LOCAL_WEAVIATE_GRPC_PORT", 50051)
-        )   
+        self.save_location_local_weaviate_grpc_port: int = int(
+            os.getenv("LOCAL_WEAVIATE_GRPC_PORT", 50051)
+        )
         self.save_location_client_manager: ClientManager = ClientManager(
             wcd_url=self.save_location_wcd_url,
             wcd_api_key=self.save_location_wcd_api_key,
@@ -199,13 +200,19 @@ async def configure(self, **kwargs):
             self.save_location_wcd_api_key = kwargs["save_location_wcd_api_key"]
             reload_client_manager = True
         if "save_location_weaviate_is_local" in kwargs:
-            self.save_location_weaviate_is_local = kwargs["save_location_weaviate_is_local"]
+            self.save_location_weaviate_is_local = kwargs[
+                "save_location_weaviate_is_local"
+            ]
             reload_client_manager = True
         if "save_location_local_weaviate_port" in kwargs:
-            self.save_location_local_weaviate_port = kwargs["save_location_local_weaviate_port"]
+            self.save_location_local_weaviate_port = kwargs[
+                "save_location_local_weaviate_port"
+            ]
             reload_client_manager = True
         if "save_location_local_weaviate_grpc_port" in kwargs:
-            self.save_location_local_weaviate_grpc_port = kwargs["save_location_local_weaviate_grpc_port"]
+            self.save_location_local_weaviate_grpc_port = kwargs[
+                "save_location_local_weaviate_grpc_port"
+            ]
             reload_client_manager = True
         if "save_trees_to_weaviate" in kwargs:
             self.config["save_trees_to_weaviate"] = kwargs["save_trees_to_weaviate"]

elysia/config.py

@@ -115,6 +115,7 @@ def is_api_key(key: str) -> bool:
         or key.lower().endswith("_access_key_id")
         or key.lower().endswith("_secret_access_key")
         or key.lower().endswith("_region_name")
+        or key.lower().endswith("_token")
     )
 
 
@@ -346,6 +347,7 @@ def configure(
                 - local_weaviate_port (int): The port to use for the local Weaviate cluster.
                 - local_weaviate_grpc_port (int): The gRPC port to use for the local Weaviate cluster.
                 - logging_level (str): The logging level to use. e.g. "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"
+                - api_keys (dict): A dictionary of API keys to set. E.g. `{"openai_apikey": "..."}`.
                 - use_feedback (bool): EXPERIMENTAL. Whether to use feedback from previous runs of the tree.
                     If True, the tree will use TrainingUpdate objects that have been saved in previous runs of the decision tree.
                     These are implemented via few-shot examples for the decision node.
@@ -627,19 +629,58 @@ def __enter__(self):
         # save existing env
         self.existing_env = deepcopy(os.environ)
 
-        # blank out the environ except for certain values
-        os.environ = {
-            "MODEL_API_BASE": self.settings.MODEL_API_BASE,
-            "WCD_URL": self.settings.WCD_URL,
-            "WCD_API_KEY": self.settings.WCD_API_KEY,
-            "WEAVIATE_IS_LOCAL": str(self.settings.WEAVIATE_IS_LOCAL),
-            "LOCAL_WEAVIATE_PORT": str(self.settings.LOCAL_WEAVIATE_PORT),
-            "LOCAL_WEAVIATE_GRPC_PORT": str(self.settings.LOCAL_WEAVIATE_GRPC_PORT),
-        }
+        items_to_remove = []
+        for key in os.environ:
+            if is_api_key(key):
+                items_to_remove.append(key)
+
+        for key in items_to_remove:
+            del os.environ[key]
+
+        if (
+            "MODEL_API_BASE" in dir(self.settings)
+            and self.settings.MODEL_API_BASE is not None
+        ):
+            os.environ["MODEL_API_BASE"] = self.settings.MODEL_API_BASE
+
+        if "WCD_URL" in dir(self.settings) and self.settings.WCD_URL is not None:
+            os.environ["WCD_URL"] = self.settings.WCD_URL
+        if (
+            "WCD_API_KEY" in dir(self.settings)
+            and self.settings.WCD_API_KEY is not None
+        ):
+            os.environ["WCD_API_KEY"] = self.settings.WCD_API_KEY
+        if (
+            "WEAVIATE_IS_LOCAL" in dir(self.settings)
+            and self.settings.WEAVIATE_IS_LOCAL is not None
+        ):
+            os.environ["WEAVIATE_IS_LOCAL"] = str(self.settings.WEAVIATE_IS_LOCAL)
+        if (
+            "LOCAL_WEAVIATE_PORT" in dir(self.settings)
+            and self.settings.LOCAL_WEAVIATE_PORT is not None
+        ):
+            os.environ["LOCAL_WEAVIATE_PORT"] = str(self.settings.LOCAL_WEAVIATE_PORT)
+        if (
+            "LOCAL_WEAVIATE_GRPC_PORT" in dir(self.settings)
+            and self.settings.LOCAL_WEAVIATE_GRPC_PORT is not None
+        ):
+            os.environ["LOCAL_WEAVIATE_GRPC_PORT"] = str(
+                self.settings.LOCAL_WEAVIATE_GRPC_PORT
+            )
 
         # update all api keys in env
+        warning_api_keys = []
         for api_key, value in self.settings.API_KEYS.items():
-            os.environ[api_key.upper()] = value
+            if isinstance(value, str):
+                os.environ[api_key.upper()] = value
+            else:
+                warning_api_keys.append(api_key)
+
+        if len(warning_api_keys) > 0:
+            self.settings.logger.warning(
+                f"The following API keys are either not found or not strings: {', '.join(warning_api_keys)}. "
+                "These have been ignored. Please ensure the API keys in the settings are available and strings."
+            )
 
         pass
 

tests/no_reqs/general/test_keymanager_env_vars_nr.py

@@ -0,0 +1,36 @@
+import os
+import pytest
+
+from elysia.config import Settings
+from fastapi.responses import JSONResponse
+import json
+from elysia.api.dependencies.common import get_user_manager
+from elysia.api.services.user import UserManager
+from weaviate.util import generate_uuid5
+
+from uuid import uuid4
+
+from elysia import tool
+from elysia.tree import Tree
+from elysia.config import ElysiaKeyManager
+
+
+def test_api_keys_are_not_passed_down():
+
+    os.environ["TEST_API_KEY"] = "1234567890"
+    settings = Settings()
+    with ElysiaKeyManager(settings):
+        assert "TEST_API_KEY" not in os.environ
+
+
+def test_modified_api_keys_are_passed_down():
+
+    os.environ["TEST_API_KEY_"] = "1234567890"
+    settings = Settings()
+    with ElysiaKeyManager(settings):
+        assert "TEST_API_KEY_" in os.environ
+
+    settings = Settings()
+    settings.configure(api_keys={"TEST_API_KEY": "1234567890"})
+    with ElysiaKeyManager(settings):
+        assert "TEST_API_KEY" in os.environ

tests/no_reqs/general/test_settings_nr.py

@@ -105,64 +105,3 @@ def test_basic_configure_local():
     assert settings.BASE_PROVIDER != global_settings.BASE_PROVIDER
 
     reset_global_settings()
-
-
-def test_model_keys():
-
-    # use settings from smart setup
-    settings = Settings()
-    tree = Tree(settings=settings)
-
-    # set the keys to wrong
-    settings.configure(
-        openrouter_api_key="wrong",
-        base_model="gemini-2.0-flash-001",
-        base_provider="openrouter/google",
-        complex_model="gemini-2.0-flash-001",
-        complex_provider="openrouter/google",
-    )
-
-    # explicitly set the adapter back to ChatAdapter
-    from dspy.adapters import ChatAdapter
-    import dspy
-
-    with dspy.context(adapter=ChatAdapter()):
-        # should error
-        with pytest.raises(Exception):
-            response, objects = tree("hi elly. use text response only")
-
-        with pytest.raises(Exception):
-            tree.create_conversation_title()
-
-        with pytest.raises(Exception):
-            tree.get_follow_up_suggestions()
-
-        # missing keys
-        settings = Settings()
-        settings.configure(
-            base_model="gemini-2.0-flash-001",
-            base_provider="openrouter/google",
-            complex_model="gemini-2.0-flash-001",
-            complex_provider="openrouter/google",
-        )
-
-        tree = Tree(settings=settings)
-
-        with pytest.raises(Exception):
-            response, objects = tree("hi elly. use text response only")
-
-        with pytest.raises(Exception):
-            tree.create_conversation_title()
-
-        with pytest.raises(Exception):
-            tree.get_follow_up_suggestions()
-
-    # now set the keys back to the .env
-    settings.configure(
-        openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
-    )
-
-    # should not error
-    response, objects = tree("hi elly. use text response only")
-    tree.create_conversation_title()
-    tree.get_follow_up_suggestions()

tests/requires_env/general/test_chunking.py

@@ -180,10 +180,7 @@ async def test_correct_vectoriser_new():
     ]
 
     # create collection with vectoriser
-    full_vectoriser = Configure.Vectors.text2vec_openai(
-        model="text-embedding-3-large",
-        dimensions=512,
-    )
+    full_vectoriser = Configure.Vectors.text2vec_openai(model="text-embedding-3-small")
     async with client_manager.connect_to_async_client() as client:
         collection_full = await client.collections.create(
             collection_name_full,

tests/requires_env/general/test_settings.py

@@ -17,3 +17,57 @@ def test_smart_setup():
     assert settings.COMPLEX_MODEL is not None
     assert settings.BASE_PROVIDER is not None
     assert settings.COMPLEX_PROVIDER is not None
+
+
+def test_model_keys():
+
+    # use settings from smart setup
+    settings = Settings()
+    tree = Tree(settings=settings)
+
+    # set the keys to wrong
+    settings.configure(
+        openrouter_api_key="wrong",
+        base_model="gemini-2.0-flash-001",
+        base_provider="openrouter/google",
+        complex_model="gemini-2.0-flash-001",
+        complex_provider="openrouter/google",
+    )
+
+    # should error
+    with pytest.raises(Exception):
+        response, objects = tree("hi elly. use text response only")
+
+    with pytest.raises(Exception):
+        tree.create_conversation_title()
+
+    with pytest.raises(Exception):
+        tree.get_follow_up_suggestions()
+
+    # missing keys
+    settings = Settings()
+    settings.configure(
+        base_model="gemini-2.0-flash-001",
+        base_provider="openrouter/google",
+        complex_model="gemini-2.0-flash-001",
+        complex_provider="openrouter/google",
+    )
+
+    tree = Tree(settings=settings)
+
+    with pytest.raises(Exception):
+        response, objects = tree("hi elly. use text response only")
+
+    with pytest.raises(Exception):
+        tree.create_conversation_title()
+
+    with pytest.raises(Exception):
+        tree.get_follow_up_suggestions()
+
+    # now set the keys back to the .env
+    settings.configure(
+        openrouter_api_key=os.getenv("OPENROUTER_API_KEY"),
+    )
+
+    # should not error
+    response, objects = tree("hi elly. use text response only")