ms-enclave 0.0.3__py3-none-any.whl

ms_enclave/sandbox/tools/tool_info.py

@@ -0,0 +1,141 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Literal, Optional, TypeAlias, Union, get_args, get_type_hints
+
+from docstring_parser import Docstring, parse
+from pydantic import BaseModel, Field
+
+from ms_enclave.utils.json_schema import JSONSchema, JSONType, json_schema, python_type_to_json_type
+
+ToolParam: TypeAlias = JSONSchema
+"""Description of tool parameter in JSON Schema format."""
+
+
+class ToolParams(BaseModel):
+    """Description of tool parameters object in JSON Schema format."""
+
+    type: Literal['object'] = Field(default='object', description="Params type (always 'object')")
+    properties: Dict[str, ToolParam] = Field(default_factory=dict, description='Tool function parameters.')
+    required: List[str] = Field(default_factory=list, description='List of required fields.')
+    additionalProperties: bool = Field(
+        default=False, description='Are additional object properties allowed? (always `False`)'
+    )
+
+
+class ToolInfo(BaseModel):
+    """Specification of a tool (JSON Schema compatible)
+
+    If you are implementing a ModelAPI, most LLM libraries can
+    be passed this object (dumped to a dict) directly as a function
+    specification. For example, in the OpenAI provider:
+
+    ```python
+    ChatCompletionToolParam(
+        type="function",
+        function=tool.model_dump(exclude_none=True),
+    )
+    ```
+
+    In some cases the field names don't match up exactly. In that case
+    call `model_dump()` on the `parameters` field. For example, in the
+    Anthropic provider:
+
+    ```python
+    ToolParam(
+        name=tool.name,
+        description=tool.description,
+        input_schema=tool.parameters.model_dump(exclude_none=True),
+    )
+    ```
+    """
+
+    name: str = Field(description='Name of tool.')
+    description: str = Field(description='Short description of tool.')
+    parameters: ToolParams = Field(default_factory=ToolParams, description='JSON Schema of tool parameters object.')
+    options: Optional[Dict[str, object]] = Field(
+        default=None,
+        description=
+        'Optional property bag that can be used by the model provider to customize the implementation of the tool'
+    )
+
+
+def parse_tool_info(func: Callable[..., Any]) -> ToolInfo:
+    # tool may already have registry attributes w/ tool info
+
+    if (getattr(func, 'name', None) and getattr(func, 'description', None) and getattr(func, 'parameters', None)):
+        return ToolInfo(
+            name=func.name,
+            description=func.description,
+            parameters=func.parameters,
+        )
+
+    signature = inspect.signature(func)
+    type_hints = get_type_hints(func)
+    docstring = inspect.getdoc(func)
+    parsed_docstring: Optional[Docstring] = parse(docstring) if docstring else None
+
+    info = ToolInfo(name=func.__name__, description='')
+
+    for param_name, param in signature.parameters.items():
+        tool_param = ToolParam()
+
+        # Parse docstring
+        docstring_info = parse_docstring(docstring, param_name)
+
+        # Get type information from type annotations
+        if param_name in type_hints:
+            tool_param = json_schema(type_hints[param_name])
+        # as a fallback try to parse it from the docstring
+        # (this is minimally necessary for backwards compatiblity
+        #  with tools gen1 type parsing, which only used docstrings)
+        elif 'docstring_type' in docstring_info:
+            json_type = python_type_to_json_type(docstring_info['docstring_type'])
+            if json_type and (json_type in get_args(JSONType)):
+                tool_param = ToolParam(type=json_type)
+
+        # Get default value
+        if param.default is param.empty:
+            info.parameters.required.append(param_name)
+        else:
+            tool_param.default = param.default
+
+        # Add description from docstring
+        if 'description' in docstring_info:
+            tool_param.description = docstring_info['description']
+
+        # append the tool param
+        info.parameters.properties[param_name] = tool_param
+
+    # Add function description if available
+    if parsed_docstring:
+        if parsed_docstring.description:
+            info.description = parsed_docstring.description.strip()
+        elif parsed_docstring.long_description:
+            info.description = parsed_docstring.long_description.strip()
+        elif parsed_docstring.short_description:
+            info.description = parsed_docstring.short_description.strip()
+
+        # Add examples if available
+        if parsed_docstring.examples:
+            examples = '\n\n'.join([(example.description or '') for example in parsed_docstring.examples])
+            info.description = f'{info.description}\n\nExamples\n\n{examples}'
+
+    return info
+
+
+def parse_docstring(docstring: Optional[str], param_name: str) -> Dict[str, str]:
+    if not docstring:
+        return {}
+
+    parsed_docstring: Docstring = parse(docstring)
+
+    for param in parsed_docstring.params:
+        if param.arg_name == param_name:
+            schema: Dict[str, str] = {'description': param.description or ''}
+
+            if param.type_name:
+                schema['docstring_type'] = param.type_name
+
+            return schema
+
+    return {}

ms_enclave/utils/__init__.py

@@ -0,0 +1 @@
+from .logger import get_logger

ms_enclave/utils/json_schema.py

@@ -0,0 +1,208 @@
+import types
+import typing
+from copy import deepcopy
+from dataclasses import is_dataclass
+from datetime import date, datetime, time
+from enum import EnumMeta
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+    is_typeddict,
+)
+
+from pydantic import BaseModel, Field
+
+JSONType = Literal['string', 'integer', 'number', 'boolean', 'array', 'object', 'null']
+"""Valid types within JSON schema."""
+
+
+class JSONSchema(BaseModel):
+    """JSON Schema for type."""
+
+    type: Optional[JSONType] = Field(default=None)
+    """JSON type of tool parameter."""
+
+    format: Optional[str] = Field(default=None)
+    """Format of the parameter (e.g. date-time)."""
+
+    description: Optional[str] = Field(default=None)
+    """Parameter description."""
+
+    default: Any = Field(default=None)
+    """Default value for parameter."""
+
+    enum: Optional[List[Any]] = Field(default=None)
+    """Valid values for enum parameters."""
+
+    items: Optional['JSONSchema'] = Field(default=None)
+    """Valid type for array parameters."""
+
+    properties: Optional[Dict[str, 'JSONSchema']] = Field(default=None)
+    """Valid fields for object parametrs."""
+
+    additionalProperties: Optional[Union['JSONSchema', bool]] = Field(default=None)
+    """Are additional properties allowed?"""
+
+    anyOf: Optional[List['JSONSchema']] = Field(default=None)
+    """Valid types for union parameters."""
+
+    required: Optional[List[str]] = Field(default=None)
+    """Required fields for object parameters."""
+
+
+def json_schema(t: Type[Any]) -> JSONSchema:
+    """Provide a JSON Schema for the specified type.
+
+    Schemas can be automatically inferred for a wide variety of
+    Python class types including Pydantic BaseModel, dataclasses,
+    and typed dicts.
+
+    Args:
+        t: Python type
+
+    Returns:
+        JSON Schema for type.
+    """
+    origin = get_origin(t)
+    args = get_args(t)
+
+    if origin is None:
+        if t is int:
+            return JSONSchema(type='integer')
+        elif t is float:
+            return JSONSchema(type='number')
+        elif t is str:
+            return JSONSchema(type='string')
+        elif t is bool:
+            return JSONSchema(type='boolean')
+        elif t is datetime:
+            return JSONSchema(type='string', format='date-time')
+        elif t is date:
+            return JSONSchema(type='string', format='date')
+        elif t is time:
+            return JSONSchema(type='string', format='time')
+        elif t is list or t is set:
+            return JSONSchema(type='array', items=JSONSchema())
+        elif t is dict:
+            return JSONSchema(type='object', additionalProperties=JSONSchema())
+        elif (is_dataclass(t) or is_typeddict(t) or (isinstance(t, type) and issubclass(t, BaseModel))):
+            return cls_json_schema(t)
+        elif isinstance(t, EnumMeta):
+            return JSONSchema(enum=[item.value for item in t])
+        elif t is type(None):
+            return JSONSchema(type='null')
+        else:
+            return JSONSchema()
+    elif (origin is list or origin is List or origin is tuple or origin is Tuple or origin is set or origin is Set):
+        return JSONSchema(type='array', items=json_schema(args[0]) if args else JSONSchema())
+    elif origin is dict or origin is Dict:
+        return JSONSchema(
+            type='object',
+            additionalProperties=json_schema(args[1]) if len(args) > 1 else JSONSchema(),
+        )
+    elif origin is Union or origin is types.UnionType:
+        return JSONSchema(anyOf=[json_schema(arg) for arg in args])
+    elif origin is Optional:
+        return JSONSchema(anyOf=[json_schema(arg) for arg in args] + [JSONSchema(type='null')])
+    elif origin is typing.Literal:
+        return JSONSchema(enum=list(args))
+
+    return JSONSchema()  # Default case if we can't determine the type
+
+
+def cls_json_schema(cls: Type[Any]) -> JSONSchema:
+    properties: Dict[str, JSONSchema] = {}
+    required: List[str] = []
+
+    if is_dataclass(cls):
+        fields = cls.__dataclass_fields__  # type: ignore
+        for name, field in fields.items():
+            properties[name] = json_schema(field.type)  # type: ignore
+            if field.default == field.default_factory:
+                required.append(name)
+    elif isinstance(cls, type) and issubclass(cls, BaseModel):
+        schema = cls.model_json_schema()
+        schema = resolve_schema_references(schema)
+        for name, prop in schema.get('properties', {}).items():
+            properties[name] = JSONSchema(**prop)
+        required = schema.get('required', [])
+    elif is_typeddict(cls):
+        annotations = get_type_hints(cls)
+        for name, type_hint in annotations.items():
+            properties[name] = json_schema(type_hint)
+            if name in cls.__required_keys__:
+                required.append(name)
+
+    return JSONSchema(
+        type='object',
+        properties=properties,
+        required=required if required else None,
+        additionalProperties=False,
+    )
+
+
+def python_type_to_json_type(python_type: Optional[str]) -> JSONType:
+    if python_type == 'str':
+        return 'string'
+    elif python_type == 'int':
+        return 'integer'
+    elif python_type == 'float':
+        return 'number'
+    elif python_type == 'bool':
+        return 'boolean'
+    elif python_type == 'list':
+        return 'array'
+    elif python_type == 'dict':
+        return 'object'
+    elif python_type == 'None':
+        return 'null'
+    elif python_type is None:
+        # treat 'unknown' as string as anything can be converted to string
+        return 'string'
+    else:
+        raise ValueError(f'Unsupported type: {python_type} for Python to JSON conversion.')
+
+
+def resolve_schema_references(schema: Dict[str, Any]) -> Dict[str, Any]:
+    """Resolves all $ref references in a JSON schema by inlining the definitions."""
+    schema = deepcopy(schema)
+    definitions = schema.pop('$defs', {})
+
+    def _resolve_refs(obj: Any) -> Any:
+        if isinstance(obj, dict):
+            if '$ref' in obj and obj['$ref'].startswith('#/$defs/'):
+                ref_key = obj['$ref'].split('/')[-1]
+                if ref_key in definitions:
+                    # Replace with a deep copy of the definition
+                    resolved = deepcopy(definitions[ref_key])
+                    # Process any nested references in the definition
+                    resolved = _resolve_refs(resolved)
+
+                    # Merge in the current object fields, which should take priority
+                    # This means that if you have e.g.
+                    # {"$ref": "#/$defs/SubType", "description": "subtype of type SubType"},
+                    # and SubType resolves to
+                    # {"description": "The SubType Class", "parameters": {"param1": {"type": "string"}}},
+                    # the final result will be:
+                    # {"description": "subtype of type SubType", "parameters": {"param1": {"type": "string"}}}
+                    return resolved | {k: o for k, o in obj.items() if k != '$ref'}
+
+            # Process all entries in the dictionary
+            return {k: _resolve_refs(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [_resolve_refs(item) for item in obj]
+        else:
+            return obj
+
+    return cast(Dict[str, Any], _resolve_refs(schema))

ms_enclave/utils/logger.py

@@ -0,0 +1,170 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import importlib.util
+import logging
+import os
+import sys
+import threading
+from types import MethodType
+from typing import Optional
+
+init_loggers = {}
+
+# ANSI color helpers for levelname coloring in TTY streams
+RESET = '\033[0m'
+LEVEL_COLORS = {
+    'DEBUG': '\033[34m',  # Blue
+    'INFO': '\033[32m',  # Green
+    'WARNING': '\033[33m',  # Yellow
+    'ERROR': '\033[31m',  # Red
+    'CRITICAL': '\033[35m',  # Magenta
+}
+
+logger_format = logging.Formatter('[%(levelname)s:%(name)s] %(message)s')
+
+info_set = set()
+warning_set = set()
+_once_lock = threading.Lock()
+
+
+class ColorFormatter(logging.Formatter):
+    """Formatter that colors only the levelname for TTY streams."""
+
+    def __init__(self, fmt: str, datefmt: Optional[str] = None, style: str = '%', use_color: bool = True) -> None:
+        super().__init__(fmt=fmt, datefmt=datefmt, style=style)
+        self.use_color = use_color
+
+    def format(self, record: logging.LogRecord) -> str:
+        original_levelname = record.levelname
+        try:
+            if self.use_color:
+                color = LEVEL_COLORS.get(record.levelname, '')
+                if color:
+                    record.levelname = f'{color}{record.levelname}{RESET}'
+            return super().format(record)
+        finally:
+            record.levelname = original_levelname
+
+
+def _should_use_color(stream) -> bool:
+    """Decide if we should use colors for a given stream based on TTY and env."""
+    # Respect NO_COLOR to disable, FORCE_COLOR or LOG_COLOR=1 to force enable
+    if os.getenv('NO_COLOR'):
+        return False
+    if os.getenv('FORCE_COLOR') or os.getenv('LOG_COLOR') == '1':
+        return True
+    try:
+        return hasattr(stream, 'isatty') and stream.isatty()
+    except Exception:
+        return False
+
+
+def info_once(self: logging.Logger, msg: str, *args, **kwargs) -> None:
+    hash_id = kwargs.pop('hash_id', msg)
+    with _once_lock:
+        if hash_id in info_set:
+            return
+        info_set.add(hash_id)
+    self.info(msg, *args, **kwargs)
+
+
+def warning_once(self: logging.Logger, msg: str, *args, **kwargs) -> None:
+    hash_id = kwargs.pop('hash_id', msg)
+    with _once_lock:
+        if hash_id in warning_set:
+            return
+        warning_set.add(hash_id)
+    self.warning(msg, *args, **kwargs)
+
+
+def _update_handler_levels(logger: logging.Logger, log_level: int) -> None:
+    """Set all handler levels to the given log level."""
+    for handler in logger.handlers:
+        handler.setLevel(log_level)
+
+
+def get_logger(log_file: Optional[str] = None, log_level: Optional[int] = None, file_mode: str = 'w'):
+    """Get project logger configured with colored console output and optional file output.
+
+    Args:
+        log_file: Log filename. If specified, a FileHandler will be added to the logger.
+        log_level: Logging level. If None, resolve from env LOG_LEVEL (default INFO).
+        file_mode: Mode to open the log file if log_file is provided (default 'w').
+    """
+    if log_level is None:
+        env_level = os.getenv('LOG_LEVEL', 'INFO').upper()
+        log_level = getattr(logging, env_level, logging.INFO)
+
+    logger_name = __name__.split('.')[0]
+    logger = logging.getLogger(logger_name)
+    logger.propagate = False
+
+    # If logger is already initialized, just ensure file handler and update handler levels.
+    if logger_name in init_loggers:
+        add_file_handler_if_needed(logger, log_file, file_mode, log_level)
+        _update_handler_levels(logger, log_level)
+        return logger
+
+    # Handle duplicate logs to the console (PyTorch DDP root StreamHandler quirk)
+    for handler in logger.root.handlers:
+        if isinstance(handler, logging.StreamHandler):
+            handler.setLevel(logging.ERROR)
+
+    # Console handler with colorized levelname when appropriate
+    stream_handler = logging.StreamHandler(stream=sys.stderr)
+    use_color = _should_use_color(getattr(stream_handler, 'stream', sys.stderr))
+    color_fmt = ColorFormatter('[%(levelname)s:%(name)s] %(message)s', use_color=use_color)
+    stream_handler.setFormatter(color_fmt)
+    stream_handler.setLevel(log_level)
+    logger.addHandler(stream_handler)
+
+    # Optional file handler (no color)
+    if log_file is not None:
+        file_handler = logging.FileHandler(log_file, file_mode)
+        file_handler.setFormatter(logger_format)
+        file_handler.setLevel(log_level)
+        logger.addHandler(file_handler)
+
+    logger.setLevel(log_level)
+    init_loggers[logger_name] = True
+    logger.info_once = MethodType(info_once, logger)
+    logger.warning_once = MethodType(warning_once, logger)
+    return logger
+
+
+logger = get_logger()
+
+
+def add_file_handler_if_needed(logger: logging.Logger, log_file: Optional[str], file_mode: str, log_level: int) -> None:
+    """Attach a FileHandler for the given log_file if not already present.
+
+    Ensures:
+    - Only one FileHandler per log file path.
+    - FileHandler uses the standard, uncolored formatter.
+    - FileHandler level matches the requested log_level.
+    """
+    if log_file is None:
+        return
+
+    # Only worker 0 writes logs when torch DDP is present
+    if importlib.util.find_spec('torch') is not None:
+        is_worker0 = int(os.getenv('LOCAL_RANK', -1)) in {-1, 0}
+    else:
+        is_worker0 = True
+
+    if not is_worker0:
+        return
+
+    abs_path = os.path.abspath(log_file)
+    for handler in logger.handlers:
+        if isinstance(handler, logging.FileHandler):
+            # If a handler is already logging to the same file, just update it
+            if getattr(handler, 'baseFilename', None) == abs_path:
+                handler.setFormatter(logger_format)
+                handler.setLevel(log_level)
+                return
+
+    # Add a new file handler for this log file
+    file_handler = logging.FileHandler(abs_path, file_mode)
+    file_handler.setFormatter(logger_format)
+    file_handler.setLevel(log_level)
+    logger.addHandler(file_handler)

ms_enclave/version.py

@@ -0,0 +1,2 @@
+__version__ = '0.0.3'
+__release_date__ = '2025-11-27 12:00:00'