aixtools 0.1.6__py3-none-any.whl → 0.1.7__py3-none-any.whl

aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.6'
-__version_tuple__ = version_tuple = (0, 1, 6)
+__version__ = version = '0.1.7'
+__version_tuple__ = version_tuple = (0, 1, 7)
 
 __commit_id__ = commit_id = None

aixtools/vault/vault.py

@@ -4,7 +4,7 @@ Provides a Vault client for storing and retrieving user service api keys.
 """
 
 import logging
-from typing import Optional
+from typing import Dict, Optional
 
 import hvac
 from hvac.exceptions import InvalidPath
@@ -27,47 +27,74 @@ class VaultClient:
         if not self.client.is_authenticated():
             raise VaultAuthError("Vault client authentication failed. Check vault_token.")
 
+    def _get_secret_path(self, user_id: str, service_name: str) -> str:
+        """Generate the vault secret path for a user and service."""
+        return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+
     def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
         """
         Store user's service api key in the Vault at the specified vault mount
         point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+
+        This is a convenience method for storing a single API key.
+        For storing multiple secrets, use store_user_service_secret().
+        """
+        secret_dict = {"user-api-key": user_api_key}
+        self.store_user_service_secret(user_id=user_id, service_name=service_name, secret_data=secret_dict)
+
+    def read_user_service_api_key(self, *, user_id: str, service_name: str) -> Optional[str]:
+        """
+        Read user's service api key in from vault at the specified mount point,
+        where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+
+        This is a convenience method for reading a single API key.
+        For reading multiple secrets, use read_user_service_secret().
+        """
+        secret_data = self.read_user_service_secret(user_id=user_id, service_name=service_name)
+        if secret_data is None:
+            return None
+        return secret_data.get("user-api-key")
+
+    def store_user_service_secret(self, *, user_id: str, service_name: str, secret_data: Dict[str, str]):
+        """
+        Store complete user service secret with multiple key-value pairs in the Vault
+        at the specified vault mount point, where the path is <path_prefix>/<env>/<user_id>/<service_name>.
         """
         secret_path = None
         try:
-            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
-            print("secret_path", secret_path)
-            secret_dict = {"user-api-key": user_api_key}
+            secret_path = self._get_secret_path(user_id, service_name)
+            logger.info("Writing complete secret to path %s", secret_path)
             self.client.secrets.kv.v2.create_or_update_secret(
-                secret_path, secret=secret_dict, mount_point=config.VAULT_MOUNT_POINT
+                secret_path, secret=secret_data, mount_point=config.VAULT_MOUNT_POINT
             )
 
-            logger.info("Secret written to path %s", secret_path)
+            logger.info("Complete secret written to path %s", secret_path)
         except Exception as e:
-            logger.error("Failed to write secret to path %s: %s", secret_path, str(e))
+            logger.error("Failed to write complete secret to path %s: %s", secret_path, str(e))
             raise VaultAuthError(e) from e
 
-    def read_user_service_api_key(self, *, user_id: str, service_name) -> Optional[str]:
+    def read_user_service_secret(self, *, user_id: str, service_name: str) -> Optional[Dict[str, str]]:
         """
-        Read user's service api key in from vault at the specified mount point,
+        Read complete user service secret from vault at the specified mount point,
         where the path is <path_prefix>/<env>/<user_id>/<service_name>.
+        Returns all key-value pairs in the secret or None if the secret doesn't exist.
         """
         secret_path = None
 
         try:
-            secret_path = f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
-            logger.info("Reading secret from path %s", secret_path)
+            secret_path = self._get_secret_path(user_id, service_name)
+            logger.info("Reading complete secret from path %s", secret_path)
             response = self.client.secrets.kv.v2.read_secret_version(
                 secret_path, mount_point=config.VAULT_MOUNT_POINT, raise_on_deleted_version=True
             )
             secret_data = response["data"]["data"]
-            user_api_key = secret_data["user-api-key"]
-            logger.info("Secret read from path %s ", secret_path)
-            return user_api_key
+            logger.info("Complete secret read from path %s", secret_path)
+            return secret_data
         except InvalidPath:
             # Secret path does not exist
-            logger.warning("Secret path does not exist %s ", secret_path)
+            logger.warning("Secret path does not exist %s", secret_path)
             return None
 
         except Exception as e:
-            logger.error("Failed to read secret from path %s: %s", secret_path, str(e))
+            logger.error("Failed to read complete secret from path %s: %s", secret_path, str(e))
             raise VaultAuthError(e) from e

{aixtools-0.1.6.dist-info → aixtools-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.1.6
+Version: 0.1.7
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.1.6.dist-info → aixtools-0.1.7.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=riGXiVTWXmtdoju9hVCWvTxpszEMAAIK0sZZWoLKlnU,704
+aixtools/_version.py,sha256=szvPIs2C82UunpzuvVg3MbF4QhzbBYTsVJ8DmPfq6_E,704
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -77,7 +77,7 @@ aixtools/utils/utils.py,sha256=5911Ej1ES2NU_FKIWA3CWKhKnwgjvi1aDR2aiD6Xv3E,4880
 aixtools/utils/chainlit/cl_agent_show.py,sha256=vaRuowp4BRvhxEr5hw0zHEJ7iaSF_5bo_9BH7pGPPpw,4398
 aixtools/utils/chainlit/cl_utils.py,sha256=fxaxdkcZg6uHdM8uztxdPowg3a2f7VR7B26VPY4t-3c,5738
 aixtools/vault/__init__.py,sha256=fsr_NuX3GZ9WZ7dGfe0gp_5-z3URxAfwVRXw7Xyc0dU,141
-aixtools/vault/vault.py,sha256=JeAnRnsXm_etkTUntE9ABIiNlI8M81E2b85IQHfGehI,2817
+aixtools/vault/vault.py,sha256=WkzBTEYM-Vqjyoa5x5imbEDi0ePBklMjD_aAdvIK-34,4293
 docker/mcp-base/Dockerfile,sha256=sSpbt0sasSBHHeGwPIpJpiEQMU5HGeXzerK8biVSt7Q,1547
 notebooks/example_faulty_mcp_server.ipynb,sha256=b2Cy3GXfj-gOBZ7SoUzj25F1rxp5u-32EWPHWQ-sxn8,1729
 notebooks/example_mcp_server_stdio.ipynb,sha256=ya4dRKNFU2vQxob-uIhKHGAzINXGQ6MehgKVmSCpHLk,1634
@@ -112,9 +112,9 @@ tests/unit/server/test_utils.py,sha256=kvhzdgNfsJl5tqcRBWg2yTR5GPpyrFCOmEIOuHb39
 tests/unit/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/unit/utils/test_files.py,sha256=AKFmXQqXstyKd2PreE4EmQyhQYeqOmu1Sp80MwHrf_Q,5782
 tests/unit/vault/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/unit/vault/test_vault.py,sha256=YzmJ-JZo8OLGNcmc845knCh3zp8jhkANVgKT09cuhSU,4250
-aixtools-0.1.6.dist-info/METADATA,sha256=6JFv550ISVltjdNxLIcsM0TNd-O7KS-g2ZsvCuanuJ0,18569
-aixtools-0.1.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.1.6.dist-info/entry_points.txt,sha256=dHoutULEZx7xXSqJrZdViSVjfInJibfLibi2nRXL3SE,56
-aixtools-0.1.6.dist-info/top_level.txt,sha256=ee4eF-0pqu45zCUVml0mWIhnXQgqMQper2-49BBVHLY,40
-aixtools-0.1.6.dist-info/RECORD,,
+tests/unit/vault/test_vault.py,sha256=R_RTDsralvE0JhuIrHQTb85gU9ipDJD7aoYIMdcg0o4,7264
+aixtools-0.1.7.dist-info/METADATA,sha256=PRE6AkiaiOoNeUWBs9PhV2O0jhSRRkPvqOvfodz57Lo,18569
+aixtools-0.1.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.1.7.dist-info/entry_points.txt,sha256=dHoutULEZx7xXSqJrZdViSVjfInJibfLibi2nRXL3SE,56
+aixtools-0.1.7.dist-info/top_level.txt,sha256=ee4eF-0pqu45zCUVml0mWIhnXQgqMQper2-49BBVHLY,40
+aixtools-0.1.7.dist-info/RECORD,,

tests/unit/vault/test_vault.py

@@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 from hvac.exceptions import InvalidPath
 
-from aixtools.vault.vault import VaultClient, VaultAuthError
+from aixtools.vault.vault import VaultAuthError, VaultClient
 
 
 @pytest.fixture
@@ -25,24 +25,27 @@ def valid_params():
         "env": "dev",
         "user_id": "test-user",
         "service_name": "test-service",
-        "user_api_key": "test-api-key"
+        "user_api_key": "test-api-key",
     }
 
 
 def test_store_user_service_api_key_success(patched_vault_client, valid_params):
     patched_vault_client.client.secrets.kv.v2.create_or_update_secret.return_value = {}
 
-    patched_vault_client.store_user_service_api_key(user_id=valid_params["user_id"],
-                                                    service_name=valid_params['service_name'],
-                                                    user_api_key=valid_params['user_api_key'])
+    patched_vault_client.store_user_service_api_key(
+        user_id=valid_params["user_id"],
+        service_name=valid_params["service_name"],
+        user_api_key=valid_params["user_api_key"],
+    )
 
-    secret_path = f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
+    secret_path = (
+        f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
+    )
 
-    print("expected secret path", secret_path)
     patched_vault_client.client.secrets.kv.v2.create_or_update_secret.assert_called_once_with(
         secret_path,
         secret={"user-api-key": valid_params["user_api_key"]},
-        mount_point=valid_params["vault_mount_point"]
+        mount_point=valid_params["vault_mount_point"],
     )
 
 
@@ -50,33 +53,24 @@ def test_store_user_service_api_key_invalid_path(patched_vault_client, valid_par
     patched_vault_client.client.secrets.kv.v2.create_or_update_secret.side_effect = Exception("Invalid Path")
 
     with pytest.raises(VaultAuthError, match="Invalid Path"):
-        patched_vault_client.store_user_service_api_key(user_id=valid_params["user_id"],
-                                                    service_name=valid_params['service_name'],
-                                                    user_api_key=valid_params['user_api_key'])
+        patched_vault_client.store_user_service_api_key(
+            user_id=valid_params["user_id"],
+            service_name=valid_params["service_name"],
+            user_api_key=valid_params["user_api_key"],
+        )
         patched_vault_client.client.assert_not_called()
 
 
 def test_read_user_service_api_key_success(patched_vault_client):
     """Test successful read of user service API key."""
-    mock_response = {
-        "data": {
-            "data": {
-                "user-api-key": "test-api-key"
-            }
-        }
-    }
+    mock_response = {"data": {"data": {"user-api-key": "test-api-key"}}}
     patched_vault_client.client.secrets.kv.v2.read_secret_version.return_value = mock_response
 
-    result = patched_vault_client.read_user_service_api_key(
-        user_id="test-user",
-        service_name="test-service"
-    )
+    result = patched_vault_client.read_user_service_api_key(user_id="test-user", service_name="test-service")
 
     assert result == "test-api-key"
     patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
-        "path/dev/test-user/test-service",
-        mount_point='secret',
-        raise_on_deleted_version=True
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
     )
 
 
@@ -84,16 +78,11 @@ def test_read_user_service_api_key_secret_not_found(patched_vault_client):
     """Test read_user_service_api_key when the secret path does not exist."""
     patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = InvalidPath
 
-    result = patched_vault_client.read_user_service_api_key(
-        user_id="test-user",
-        service_name="test-service"
-    )
+    result = patched_vault_client.read_user_service_api_key(user_id="test-user", service_name="test-service")
 
     assert result is None
     patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
-        "path/dev/test-user/test-service",
-        mount_point="secret",
-        raise_on_deleted_version=True
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
     )
 
 
@@ -102,13 +91,81 @@ def test_read_user_service_api_key_unexpected_error(patched_vault_client):
     patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Unexpected error")
 
     with pytest.raises(VaultAuthError, match="Unexpected error"):
-        patched_vault_client.read_user_service_api_key(
-            user_id="test-user",
-            service_name="test-service"
+        patched_vault_client.read_user_service_api_key(user_id="test-user", service_name="test-service")
+
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_store_user_service_secret_success(patched_vault_client, valid_params):
+    """Test successful storage of complete user service secret."""
+    secret_data = {"api_key": "test-api-key", "token": "test-token", "endpoint": "https://api.example.com"}
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.return_value = {}
+
+    patched_vault_client.store_user_service_secret(
+        user_id=valid_params["user_id"],
+        service_name=valid_params["service_name"],
+        secret_data=secret_data,
+    )
+
+    secret_path = (
+        f"{valid_params['path_prefix']}/{valid_params['env']}/{valid_params['user_id']}/{valid_params['service_name']}"
+    )
+
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.assert_called_once_with(
+        secret_path,
+        secret=secret_data,
+        mount_point=valid_params["vault_mount_point"],
+    )
+
+
+def test_store_user_service_secret_error(patched_vault_client, valid_params):
+    """Test store_user_service_secret when an error occurs."""
+    secret_data = {"api_key": "test-api-key", "token": "test-token"}
+    patched_vault_client.client.secrets.kv.v2.create_or_update_secret.side_effect = Exception("Storage error")
+
+    with pytest.raises(VaultAuthError, match="Storage error"):
+        patched_vault_client.store_user_service_secret(
+            user_id=valid_params["user_id"],
+            service_name=valid_params["service_name"],
+            secret_data=secret_data,
         )
 
+
+def test_read_user_service_secret_success(patched_vault_client):
+    """Test successful read of complete user service secret."""
+    secret_data = {"api_key": "test-api-key", "token": "test-token", "endpoint": "https://api.example.com"}
+    mock_response = {"data": {"data": secret_data}}
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.return_value = mock_response
+
+    result = patched_vault_client.read_user_service_secret(user_id="test-user", service_name="test-service")
+
+    assert result == secret_data
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_secret_not_found(patched_vault_client):
+    """Test read_user_service_secret when the secret path does not exist."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = InvalidPath
+
+    result = patched_vault_client.read_user_service_secret(user_id="test-user", service_name="test-service")
+
+    assert result is None
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
+    )
+
+
+def test_read_user_service_secret_error(patched_vault_client):
+    """Test read_user_service_secret when an unexpected exception occurs."""
+    patched_vault_client.client.secrets.kv.v2.read_secret_version.side_effect = Exception("Read error")
+
+    with pytest.raises(VaultAuthError, match="Read error"):
+        patched_vault_client.read_user_service_secret(user_id="test-user", service_name="test-service")
+
     patched_vault_client.client.secrets.kv.v2.read_secret_version.assert_called_once_with(
-        "path/dev/test-user/test-service",
-        mount_point="secret",
-        raise_on_deleted_version=True
+        "path/dev/test-user/test-service", mount_point="secret", raise_on_deleted_version=True
     )