aixtools 0.1.8.post1__tar.gz → 0.1.9__tar.gz

{aixtools-0.1.8.post1 → aixtools-0.1.9}/PKG-INFO

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

{aixtools-0.1.8.post1 → aixtools-0.1.9}/aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.8.post1'
-__version_tuple__ = version_tuple = (0, 1, 8, 'post1')
+__version__ = version = '0.1.9'
+__version_tuple__ = version_tuple = (0, 1, 9)
 
-__commit_id__ = commit_id = 'g6ff4d2790'
+__commit_id__ = commit_id = 'g3e272f9d5'

aixtools-0.1.9/aixtools/compliance/private_data.py

@@ -0,0 +1,136 @@
+import json
+from pathlib import Path
+
+from fastmcp import Context
+
+from aixtools.server.path import get_workspace_path
+
+PRIVATE_DATA_FILE = ".private_data"
+
+
+class PrivateData:
+    """
+    Class to manage private data file in the workspace.
+
+    The information is stored in a JSON file named `.private_data` within the workspace directory.
+    If the file does not exist, it indicates that there is no private data.
+
+    IMPORTANT: All modifications save the data to the file immediately.
+
+    FIXME: We should add some level of mutex/locking to prevent concurrent writes.
+    """
+
+    def __init__(self, ctx: Context | None = None):
+        self.ctx: Context | None = ctx
+        self._has_private_data: bool = False  # Flag indicating if private data exists
+        self._private_datasets: list[str] = []  # List of private datasets
+        self._idap_datasets: list[str] = []  # List of dataset with IDAP
+        self.load()
+
+    def add_private_dataset(self, dataset_name: str) -> None:
+        """
+        Add a private dataset to the list.
+        Save the state after modification.
+        """
+        if dataset_name not in self._private_datasets:
+            self._private_datasets.append(dataset_name)
+            self._has_private_data = True
+            self.save()
+
+    def add_idap_dataset(self, dataset_name: str) -> None:
+        """
+        Add a dataset with IDAP to the list.
+        This also adds it to the private datasets if not already present.
+        Save the state after modification.
+        """
+        if not self.has_idap_dataset(dataset_name):
+            self._idap_datasets.append(dataset_name)
+            self._has_private_data = True
+            # An IDAP dataset is also a private dataset
+            if not self.has_private_dataset(dataset_name):
+                self._private_datasets.append(dataset_name)
+            self.save()
+
+    def get_private_datasets(self) -> list[str]:
+        """Get the list of private datasets as a copy (to avoid modification)."""
+        return list(self._private_datasets)
+
+    def get_idap_datasets(self) -> list[str]:
+        """Get the list of datasets with IDAP as a copy (to avoid modification)."""
+        return list(self._idap_datasets)
+
+    def has_private_dataset(self, dataset_name: str) -> bool:
+        """Check if a specific private dataset exists."""
+        return dataset_name in self._private_datasets
+
+    def has_idap_dataset(self, dataset_name: str) -> bool:
+        """Check if a specific dataset with IDAP exists."""
+        return dataset_name in self._idap_datasets
+
+    @property
+    def has_private_data(self) -> bool:
+        """Check if private data exists."""
+        return self._has_private_data
+
+    @has_private_data.setter
+    def has_private_data(self, value: bool) -> None:
+        """
+        Set the flag indicating if private data exists.
+        Save the state after modification.
+        """
+        self._has_private_data = value
+        if not value:
+            self._private_datasets = []
+            self._idap_datasets = []
+        self.save()
+
+    def _get_private_data_path(self) -> Path:
+        """Get the path to the private data file in the workspace."""
+        return get_workspace_path(service_name=None, ctx=self.ctx) / PRIVATE_DATA_FILE
+
+    def _has_private_data_file(self) -> bool:
+        """Check if the private data file exists in the workspace."""
+        private_data_path = self.get_private_data_path()  # type: ignore
+        return private_data_path.exists()
+
+    def save(self) -> None:
+        """Save content to the private data file in the workspace."""
+        private_data_path = self._get_private_data_path()
+        # No private data? Delete the file if it exists
+        if not self.has_private_data:
+            private_data_path.unlink(missing_ok=True)
+            return
+        # If there is private data, serialize this object as JSON
+        private_data_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(private_data_path, "w") as f:
+            # Dump class as JSON, excluding the context
+            data_dict = self.__dict__.copy()
+            data_dict["ctx"] = None
+            json_data = json.dumps(data_dict, indent=4)
+            f.write(json_data)
+
+    def load(self) -> None:
+        """Load content from the private data file in the workspace."""
+        private_data_path = self._get_private_data_path()
+        if not private_data_path.exists():
+            # No private data file
+            self.has_private_data = False
+            self._private_datasets = []
+            self._idap_datasets = []
+            return
+        with open(private_data_path, "r") as f:
+            data = json.load(f)
+            self.has_private_data = data.get("_has_private_data", False)
+            self._private_datasets = data.get("_private_datasets", [])
+            self._idap_datasets = data.get("_idap_datasets", [])
+
+    def __repr__(self) -> str:
+        return (
+            f"PrivateData(has_private_data={self.has_private_data}, "
+            f"private_datasets={self._private_datasets}, "
+            f"idap_datasets={self._idap_datasets}), "
+            f"file_path={self._get_private_data_path()})"
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()

{aixtools-0.1.8.post1 → aixtools-0.1.9}/aixtools/server/__init__.py

@@ -13,6 +13,10 @@ from .utils import (
     get_session_id_tuple,
     run_in_thread,
 )
+from .workspace_privacy import (
+    is_session_private,
+    set_session_private,
+)
 
 __all__ = [
     "get_workspace_path",
@@ -20,4 +24,6 @@ __all__ = [
     "container_to_host_path",
     "host_to_container_path",
     "run_in_thread",
+    "is_session_private",
+    "set_session_private",
 ]

aixtools-0.1.9/aixtools/server/workspace_privacy.py

@@ -0,0 +1,65 @@
+"""
+Workspace privacy utilities for managing session-level privacy flags.
+"""
+
+from pathlib import Path
+
+from fastmcp import Context
+
+from aixtools.logging.logging_config import get_logger
+from aixtools.server.path import get_workspace_path
+
+logger = get_logger(__name__)
+
+PRIVACY_FLAG_FILENAME = ".private_data_indicator"
+
+
+def set_session_private(ctx: Context | tuple[str, str] | None = None) -> bool:
+    """
+    Set the current session as private by creating a privacy flag file.
+
+    Creates an empty file in the session workspace directory
+    and sets it as read-only to prevent accidental removal
+
+    Args:
+        ctx: FastMCP context for user/session identification.
+             If None, uses current FastMCP request context from HTTP headers.
+             If tuple, first part is a user id (username), second part is session id (aka conversation id)
+
+    Returns:
+        bool: True if privacy flag was successfully created, False otherwise.
+    """
+    try:
+        workspace_path = Path(get_workspace_path(ctx=ctx))
+        privacy_file = workspace_path / PRIVACY_FLAG_FILENAME
+        workspace_path.mkdir(parents=True, exist_ok=True)
+        privacy_file.touch(exist_ok=True)
+        privacy_file.chmod(0o444)
+        logger.warning("Session marked as private")
+        return True
+    except (OSError, ValueError, RuntimeError) as e:
+        logger.error("Set current session as private: %s", str(e))
+        return False
+
+
+def is_session_private(ctx: Context | tuple[str, str] | None = None) -> bool:
+    """
+    Check if the current session is marked as private.
+
+    Args:
+        ctx: FastMCP context for user/session identification.
+             If None, uses current FastMCP request context from HTTP headers.
+             If tuple, first part is a user id (username), second part is session id (aka conversation id)
+
+    Returns:
+        bool: True if workspace is private (flag file exists), False otherwise.
+    """
+    try:
+        workspace_path = Path(get_workspace_path(ctx=ctx))
+        privacy_file = workspace_path / PRIVACY_FLAG_FILENAME
+        is_private = privacy_file.exists()
+        logger.info("Session privacy check, is private: %s", str(is_private))
+        return is_private
+    except (OSError, ValueError, RuntimeError) as e:
+        logger.error("Check session privacy: %s, assuming not private!", str(e))
+        return False

{aixtools-0.1.8.post1 → aixtools-0.1.9}/aixtools/vault/vault.py

@@ -27,9 +27,11 @@ 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 _get_secret_path(self, user_id: str, service_name: Optional[str] = None) -> str:
+        """Generate the vault secret path for a user and optionally a service."""
+        if service_name:
+            return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}/{service_name}"
+        return f"{config.VAULT_PATH_PREFIX}/{config.VAULT_ENV}/{user_id}"
 
     def store_user_service_api_key(self, *, user_id: str, service_name: str, user_api_key: str):
         """
@@ -98,3 +100,38 @@ class VaultClient:
         except Exception as e:
             logger.error("Failed to read complete secret from path %s: %s", secret_path, str(e))
             raise VaultAuthError(e) from e
+
+    def list_user_secret_keys(self, *, user_id: str) -> list[str]:
+        """
+        List all secret keys (service names) for a user, optionally filtered by service name.
+
+        Args:
+            user_id: The user ID to list secrets for
+            service_name: Optional service name to filter results. If provided, returns only this service if it exists.
+
+        Returns:
+            List of service names (secret keys) for the user. Empty list if no secrets exist.
+        """
+        try:
+            # List all services for user
+            user_path = self._get_secret_path(user_id)
+            logger.info("Listing secret keys for user at path %s", user_path)
+
+            response = self.client.secrets.kv.v2.list_secrets(path=user_path, mount_point=config.VAULT_MOUNT_POINT)
+
+            if response and "data" in response and "keys" in response["data"]:
+                secret_keys = response["data"]["keys"]
+                # Remove trailing slashes from directory names if any
+                secret_keys = [key.rstrip("/") for key in secret_keys]
+                logger.info("Found %d secret keys for user %s", len(secret_keys), user_id)
+                return secret_keys
+            logger.info("No secret keys found for user %s", user_id)
+            return []
+
+        except InvalidPath:
+            # User path does not exist
+            logger.warning("User path does not exist for user %s", user_id)
+            return []
+        except Exception as e:
+            logger.error("Failed to list secret keys for user %s: %s", user_id, str(e))
+            raise VaultAuthError(e) from e

{aixtools-0.1.8.post1 → aixtools-0.1.9}/aixtools.egg-info/PKG-INFO

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

{aixtools-0.1.8.post1 → aixtools-0.1.9}/aixtools.egg-info/SOURCES.txt

@@ -47,6 +47,7 @@ aixtools/agents/__init__.py
 aixtools/agents/agent.py
 aixtools/agents/agent_batch.py
 aixtools/agents/prompt.py
+aixtools/compliance/private_data.py
 aixtools/db/__init__.py
 aixtools/db/database.py
 aixtools/db/vector_db.py
@@ -78,6 +79,7 @@ aixtools/server/__init__.py
 aixtools/server/app_mounter.py
 aixtools/server/path.py
 aixtools/server/utils.py
+aixtools/server/workspace_privacy.py
 aixtools/testing/__init__.py
 aixtools/testing/aix_test_model.py
 aixtools/testing/mock_tool.py
@@ -97,6 +99,7 @@ aixtools/utils/chainlit/cl_utils.py
 aixtools/vault/__init__.py
 aixtools/vault/vault.py
 docker/mcp-base/Dockerfile
+docker/mcp-base/zscaler.crt
 notebooks/example_faulty_mcp_server.ipynb
 notebooks/example_mcp_server_stdio.ipynb
 notebooks/example_raw_mcp_client.ipynb
@@ -120,6 +123,7 @@ tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_agent_executor.py
 tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py
 tests/unit/agents/__init__.py
 tests/unit/agents/test_prompt.py
+tests/unit/compliance/test_private_data.py
 tests/unit/google/__init__.py
 tests/unit/google/test_client.py
 tests/unit/mcp/__init__.py

{aixtools-0.1.8.post1 → aixtools-0.1.9}/docker/mcp-base/Dockerfile

@@ -1,17 +1,14 @@
 FROM ubuntu:22.04
 
 RUN apt-get -y update && \
-    apt-get -y install git curl sudo ca-certificates
+    apt-get -y install ca-certificates curl gcc git libcap2-bin sudo
+RUN mv /usr/bin/sudo /usr/sbin
 
-# Build argument for custom certificate file (optional)
-ARG CUSTOM_CERT_FILE
-# Copy custom certificate if specified
-RUN if [ -n "$CUSTOM_CERT_FILE" ] && [ -f "$CUSTOM_CERT_FILE" ]; then \
-    cp "$CUSTOM_CERT_FILE" /usr/local/share/ca-certificates/ && \
-    update-ca-certificates && \
-    export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt && \
-    export CURL_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt; \
-    fi
+# Add Zscaler CA certificate
+COPY ./zscaler.crt /usr/local/share/ca-certificates/
+RUN update-ca-certificates
+ENV SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt
+ENV CURL_CA_BUNDLE=/etc/ssl/certs/ca-certificates.crt
 
 # Install `uv` Python package manager
 RUN bash -o pipefail -c "curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh"

aixtools-0.1.9/docker/mcp-base/zscaler.crt

@@ -0,0 +1,28 @@
+-----BEGIN CERTIFICATE-----
+MIIE0zCCA7ugAwIBAgIJANu+mC2Jt3uTMA0GCSqGSIb3DQEBCwUAMIGhMQswCQYD
+VQQGEwJVUzETMBEGA1UECBMKQ2FsaWZvcm5pYTERMA8GA1UEBxMIU2FuIEpvc2Ux
+FTATBgNVBAoTDFpzY2FsZXIgSW5jLjEVMBMGA1UECxMMWnNjYWxlciBJbmMuMRgw
+FgYDVQQDEw9ac2NhbGVyIFJvb3QgQ0ExIjAgBgkqhkiG9w0BCQEWE3N1cHBvcnRA
+enNjYWxlci5jb20wHhcNMTQxMjE5MDAyNzU1WhcNNDIwNTA2MDAyNzU1WjCBoTEL
+MAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExETAPBgNVBAcTCFNhbiBK
+b3NlMRUwEwYDVQQKEwxac2NhbGVyIEluYy4xFTATBgNVBAsTDFpzY2FsZXIgSW5j
+LjEYMBYGA1UEAxMPWnNjYWxlciBSb290IENBMSIwIAYJKoZIhvcNAQkBFhNzdXBw
+b3J0QHpzY2FsZXIuY29tMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA
+qT7STSxZRTgEFFf6doHajSc1vk5jmzmM6BWuOo044EsaTc9eVEV/HjH/1DWzZtcr
+fTj+ni205apMTlKBW3UYR+lyLHQ9FoZiDXYXK8poKSV5+Tm0Vls/5Kb8mkhVVqv7
+LgYEmvEY7HPY+i1nEGZCa46ZXCOohJ0mBEtB9JVlpDIO+nN0hUMAYYdZ1KZWCMNf
+5J/aTZiShsorN2A38iSOhdd+mcRM4iNL3gsLu99XhKnRqKoHeH83lVdfu1XBeoQz
+z5V6gA3kbRvhDwoIlTBeMa5l4yRdJAfdpkbFzqiwSgNdhbxTHnYYorDzKfr2rEFM
+dsMU0DHdeAZf711+1CunuQIDAQABo4IBCjCCAQYwHQYDVR0OBBYEFLm33UrNww4M
+hp1d3+wcBGnFTpjfMIHWBgNVHSMEgc4wgcuAFLm33UrNww4Mhp1d3+wcBGnFTpjf
+oYGnpIGkMIGhMQswCQYDVQQGEwJVUzETMBEGA1UECBMKQ2FsaWZvcm5pYTERMA8G
+A1UEBxMIU2FuIEpvc2UxFTATBgNVBAoTDFpzY2FsZXIgSW5jLjEVMBMGA1UECxMM
+WnNjYWxlciBJbmMuMRgwFgYDVQQDEw9ac2NhbGVyIFJvb3QgQ0ExIjAgBgkqhkiG
+9w0BCQEWE3N1cHBvcnRAenNjYWxlci5jb22CCQDbvpgtibd7kzAMBgNVHRMEBTAD
+AQH/MA0GCSqGSIb3DQEBCwUAA4IBAQAw0NdJh8w3NsJu4KHuVZUrmZgIohnTm0j+
+RTmYQ9IKA/pvxAcA6K1i/LO+Bt+tCX+C0yxqB8qzuo+4vAzoY5JEBhyhBhf1uK+P
+/WVWFZN/+hTgpSbZgzUEnWQG2gOVd24msex+0Sr7hyr9vn6OueH+jj+vCMiAm5+u
+kd7lLvJsBu3AO3jGWVLyPkS3i6Gf+rwAp1OsRrv3WnbkYcFf9xjuaf4z0hRCrLN2
+xFNjavxrHmsH8jPHVvgc1VD0Opja0l/BRVauTrUaoW6tE+wFG5rEcPGS80jjHK4S
+pB5iDj2mUZH1T8lzYtuZy0ZPirxmtsk3135+CKNa2OCAhhFjE0xd
+-----END CERTIFICATE-----

{aixtools-0.1.8.post1 → aixtools-0.1.9}/scripts/test.sh

@@ -15,9 +15,16 @@ touch .env
 # Add the project root to PYTHONPATH so imports work correctly
 export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
 
-# Run all tests using pytest with coverage
-echo "Running tests with coverage..."
-pytest tests/ -v --cov=aixtools --cov-report=term-missing
-
-# Generate HTML coverage report
-# pytest tests/ --cov=aixtools --cov-report=html
+# Check if a specific test is provided as parameter
+if [ $# -eq 1 ]; then
+    # Run specific test without coverage
+    echo "Running specific test: $1"
+    pytest "$1" -v
+else
+    # Run all tests using pytest with coverage
+    echo "Running tests with coverage..."
+    pytest tests/ -v --cov=aixtools --cov-report=term-missing
+    
+    # Generate HTML coverage report
+    # pytest tests/ --cov=aixtools --cov-report=html
+fi