aixtools 0.0.0__py3-none-any.whl

aixtools/utils/utils.py

@@ -0,0 +1,167 @@
+"""
+General utility functions for string manipulation, logging, and data handling.
+"""
+
+import json
+import uuid
+from datetime import datetime
+from pathlib import Path
+
+import pandas as pd
+
+DF_SHOW_MAX_ROWS = 20
+
+
+def escape_newline(s, max_length: int = 300) -> str:
+    """Escape newlines in a string."""
+    s = str(s)
+    ss = "\\n".join(s.split("\n"))
+    if len(ss) <= max_length:
+        return ss
+    return "".join(ss[:max_length]) + "..."
+
+
+def escape_backticks(s) -> str:
+    """Escape backticks in a string."""
+    s = f"{s}"
+    return s.replace("`", "\\`")
+
+
+def find_file(path: Path, glob="*.pdf"):
+    """Recursively find all files matching the glob pattern in a directory"""
+    yield from path.rglob(glob)
+
+
+def prepend_all_lines(msg, prepend="\t", skip_first_line: bool = False) -> str:
+    """Prepend all lines of a message with a prepend."""
+    out = ""
+    for i, line in enumerate(str(msg).split("\n")):
+        if i == 0 and skip_first_line:
+            out += f"{line}\n"
+        else:
+            out += f"{prepend}{line}\n"
+    return out
+
+
+def remove_quotes(s):
+    """
+    Remove all quotes (including triple backticks with language specifications) surrounding a string.
+
+    This function strips the input string of all surrounding quotes, including single quotes, double quotes,
+    backticks, and triple backticks with or without language specifications. It continues to remove quotes
+    until none are left.
+
+    Args:
+        s (str): The input string potentially surrounded by quotes.
+
+    Returns:
+        str: The string with all surrounding quotes removed. If the input is None, returns None.
+    """
+    if s is None:
+        return None
+    s = s.strip()
+    while (
+        (s.startswith('"') and s.endswith('"'))
+        or (s.startswith("'") and s.endswith("'"))
+        or (s.startswith("`") and s.endswith("`"))
+        or ("```" in s)
+    ):
+        if "```" in s:
+            s = tripple_quote_strip(s)
+        else:
+            # Single quotes
+            s = s[1:-1].strip()
+        # Remove spaces
+        s = s.strip()
+    return s
+
+
+def tabit(s: str, prefix="\t|") -> str:
+    """Add a prefix to each line of a string for improved readability."""
+    return prefix + str(s).replace("\n", f"\n{prefix}")
+
+
+def to_str(data) -> str:
+    """Convert any data type to a readable string representation."""
+    # Primitive values, just use str()
+    if isinstance(data, str):
+        return f"'{str}'"
+    if isinstance(data, (bool, bytes, float, int)):
+        return str(data)
+    # Dataframes
+    if isinstance(data, pd.DataFrame):
+        if data.shape[0] > DF_SHOW_MAX_ROWS:
+            return f"Showing only the first {DF_SHOW_MAX_ROWS} rows out of {data.shape[0]}:\n" + data.head(
+                DF_SHOW_MAX_ROWS
+            ).to_markdown(index=False)
+        return data.to_markdown(index=False)
+    # Use json for list, dict, etc.
+    if isinstance(data, (dict, list, tuple)):
+        return json.dumps(data, indent=2, default=str)
+    return str(data)
+
+
+def truncate(s: str, max_len=76, ellipsis="...") -> str:
+    """Truncate a string to a maximum length, adding ellipsis if needed."""
+    s = str(s)
+    if len(s) > max_len:
+        return s[: max_len - len(ellipsis)] + ellipsis
+    return s
+
+
+def tripple_quote_strip(s):
+    """
+    Remove triple quotes from a string, including those with language specifications.
+
+    Eexamples:
+        ```sql SELECT * from table;```
+
+        ```Here is your code ```python c = a + b ``` This code will perform addition```
+
+    Args:
+        s (str): The input string potentially containing triple quotes.
+
+    Returns:
+        str: The string with triple quotes removed, if present.
+    """
+    if "```" not in s:
+        return s
+    left_pos, right_pos = len(s), s.rfind("```")
+    s_lower = s.lower()
+    pre_matched = ""
+    for lang in ["python3", "python2", "python", "bash", "json", "sql", ""]:
+        pre = f"```{lang}"
+        idx = s_lower.find(pre)
+        if idx != -1 and idx < left_pos:
+            left_pos = idx
+            pre_matched = pre
+    if left_pos < right_pos:
+        s = s[left_pos + len(pre_matched) : right_pos].strip()
+    return s
+
+
+def timestamp_with_uuid() -> str:
+    """Get a timestamp string + a UUID string (first 8 chars)."""
+    (yyy, hh, uu) = timestamp_uuid_tuple()
+    return f"{yyy}.{hh}.{uu[:8]}"
+
+
+def timestamp_uuid_tuple() -> tuple[str, str, str]:
+    """
+    Get a tuple of timestamp + a UUID: (YYYY-MM-DD, HH:MM:SS, UUID)
+    """
+    now = datetime.now()
+    return (now.strftime("%Y-%m-%d"), now.strftime("%H:%M:%S"), str(uuid.uuid4()))
+
+
+def str2bool(v: str | None) -> bool:
+    """Convert a string to a boolean value."""
+    if not v:
+        return False
+    return str(v).lower() in ("yes", "true", "on", "1")
+
+
+async def async_iter(items):
+    """Asynchronously iterate over items."""
+    for item in items:
+        yield item

aixtools/vault/__init__.py

@@ -0,0 +1,7 @@
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+from .vault import VaultClient
+
+__all__ = ["VaultClient"]

aixtools/vault/vault.py

@@ -0,0 +1,137 @@
+# ruff: noqa: PLR0913
+"""
+Provides a Vault client for storing and retrieving user service api keys.
+"""
+
+import logging
+from typing import Dict, Optional
+
+import hvac
+from hvac.exceptions import InvalidPath
+
+from aixtools.utils import config
+
+logger = logging.getLogger(__name__)
+
+
+class VaultAuthError(Exception):
+    """Exception raised for vault authentication errors."""
+
+
+class VaultClient:
+    """Vault client for storing and retrieving user service api keys."""
+
+    def __init__(self):
+        self.client = hvac.Client(url=config.VAULT_ADDRESS, token=config.VAULT_TOKEN)
+
+        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: 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):
+        """
+        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 = 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_data, mount_point=config.VAULT_MOUNT_POINT
+            )
+
+            logger.info("Complete secret written to path %s", secret_path)
+        except Exception as 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_secret(self, *, user_id: str, service_name: str) -> Optional[Dict[str, str]]:
+        """
+        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 = 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"]
+            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)
+            return None
+
+        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