hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/core/schema/generator.py

@@ -0,0 +1,742 @@
+"""Schema generator - converts Python signatures to JSON Schema."""
+
+import inspect
+import json
+import re
+from collections.abc import Callable
+from typing import Any, get_args, get_origin
+
+import yaml
+
+from hexdag.core.logging import get_logger
+from hexdag.core.types import (
+    get_annotated_metadata,
+    is_annotated_type,
+    is_dict_type,
+    is_list_type,
+    is_literal_type,
+    is_union_type,
+)
+
+logger = get_logger(__name__)
+
+
+class SchemaGenerator:
+    """Generate JSON Schema from Python callables.
+
+    This class introspects Python functions/methods to automatically generate
+    JSON Schema definitions. It supports:
+    - Basic types (str, int, float, bool)
+    - Literal types → enum
+    - Union types → anyOf
+    - List/Dict types → array/object
+    - Annotated types with Pydantic Field constraints
+    - Docstring extraction for descriptions
+
+    Examples
+    --------
+    >>> def my_func(name: str, count: int = 10):
+    ...     '''Example function.'''
+    ...     pass
+    >>> schema = SchemaGenerator.from_callable(my_func)
+    >>> schema['properties']['count']['default']
+    10
+    """
+
+    # Basic type mapping from Python to JSON Schema
+    BASIC_TYPE_MAP = {
+        str: {"type": "string"},
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        bool: {"type": "boolean"},
+        dict: {"type": "object"},
+        list: {"type": "array"},
+        None: {"type": "null"},
+        type(None): {"type": "null"},
+    }
+
+    @staticmethod
+    def from_callable(factory: Callable, format: str = "dict") -> dict | str:
+        """Generate schema from factory __call__ signature.
+
+        Resolution order:
+        1. Check for explicit ``_yaml_schema`` class attribute
+        2. Fall back to ``__call__`` signature introspection
+
+        This allows builder-pattern classes like ConditionalNode to define
+        explicit schemas for YAML/MCP usage.
+
+        Args
+        ----
+            factory: Callable (function, method, or callable class) to introspect
+            format: Output format - "dict", "yaml", or "json"
+
+        Returns
+        -------
+            dict | str: JSON Schema in requested format
+
+        Raises
+        ------
+        ValueError
+            If format is not one of: dict, yaml, json
+
+        Examples
+        --------
+        >>> def factory(name: str, count: int, enabled: bool = True):
+        ...     pass
+        >>> schema = SchemaGenerator.from_callable(factory)
+        >>> schema['properties']['enabled']['default']
+        True
+        >>> schema['required']
+        ['count']
+
+        >>> # Classes can define explicit schemas
+        >>> class MyNode:
+        ...     _yaml_schema = {"type": "object", "properties": {"foo": {"type": "string"}}}
+        >>> schema = SchemaGenerator.from_callable(MyNode)
+        >>> "foo" in schema.get("properties", {})
+        True
+        """
+        if format not in ("dict", "yaml", "json"):
+            raise ValueError(f"Invalid format: {format}. Must be one of: dict, yaml, json")
+
+        # Check for explicit _yaml_schema class attribute (for builder-pattern nodes)
+        yaml_schema = getattr(factory, "_yaml_schema", None)
+        if yaml_schema and isinstance(yaml_schema, dict):
+            logger.debug(f"Using explicit _yaml_schema for {factory}")
+            return SchemaGenerator._format_output(yaml_schema, format)
+
+        try:
+            sig = inspect.signature(factory)
+        except (ValueError, TypeError) as e:
+            logger.warning(f"Could not get signature for {factory}: {e}")
+            return SchemaGenerator._format_output({}, format)
+
+        properties = {}
+        required = []
+
+        # Extract param docs - try __call__ method first (for callable instances),
+        # then fall back to the factory itself
+        param_docs: dict[str, str] = {}
+        # We need the actual __call__ method to extract docstrings, not just callable check
+        call_method = getattr(factory, "__call__", None)  # noqa: B004
+        if call_method is not None:
+            param_docs = SchemaGenerator._extract_param_docs(call_method)
+        if not param_docs:
+            param_docs = SchemaGenerator._extract_param_docs(factory)
+
+        param_list = list(sig.parameters.items())
+        first_non_self_param = None
+        for pname, _ in param_list:
+            if pname not in ("self", "cls"):
+                first_non_self_param = pname
+                break
+
+        for param_name, param in sig.parameters.items():
+            # Skip special parameters
+            if param_name in ("self", "cls", "args", "kwargs"):
+                continue
+
+            # Skip 'name' if it's the first parameter (node factory convention)
+            if param_name == "name" and param_name == first_non_self_param:
+                continue
+
+            # Skip *args and **kwargs
+            if param.kind in (
+                inspect.Parameter.VAR_POSITIONAL,
+                inspect.Parameter.VAR_KEYWORD,
+            ):
+                continue
+
+            param_type = param.annotation
+
+            # Skip if no type annotation
+            if param_type == inspect.Parameter.empty:
+                # Default to string type
+                param_type = str
+
+            # Generate property schema
+            prop_schema = SchemaGenerator._type_to_json_schema(param_type)
+
+            if param_name in param_docs:
+                doc_text = param_docs[param_name]
+                prop_schema["description"] = doc_text
+
+                # For list[dict] types, try to extract nested structure from docstring
+                if (
+                    prop_schema.get("type") == "array"
+                    and prop_schema.get("items", {}).get("type") == "object"
+                ):
+                    nested_schema = SchemaGenerator._extract_nested_structure(doc_text)
+                    if nested_schema and "properties" in nested_schema:
+                        prop_schema["items"] = nested_schema
+
+            if param.default != inspect.Parameter.empty:
+                prop_schema["default"] = param.default
+            else:
+                # Required if no default
+                required.append(param_name)
+
+            properties[param_name] = prop_schema
+
+        schema = {
+            "type": "object",
+            "properties": properties,
+            "additionalProperties": False,
+        }
+
+        if required:
+            schema["required"] = required
+
+        return SchemaGenerator._format_output(schema, format)
+
+    @staticmethod
+    def _is_callable_type(type_hint: Any) -> bool:
+        """Check if type hint is a Callable type.
+
+        Examples
+        --------
+        >>> from collections.abc import Callable
+        >>> SchemaGenerator._is_callable_type(Callable)
+        True
+        >>> SchemaGenerator._is_callable_type(Callable[..., Any])
+        True
+        >>> SchemaGenerator._is_callable_type(str)
+        False
+        """
+        # Check for collections.abc.Callable
+        if type_hint is Callable:
+            return True
+
+        # Check for typing.Callable or parameterized Callable[..., ...]
+        origin = get_origin(type_hint)
+        if origin is Callable:
+            return True
+
+        # Check for callable origin from collections.abc
+        try:
+            from collections.abc import Callable as ABCCallable
+
+            if origin is ABCCallable:
+                return True
+        except ImportError:
+            pass
+
+        return False
+
+    @staticmethod
+    def _type_to_json_schema(type_hint: Any) -> dict:
+        """Convert Python type hint to JSON Schema type.
+
+        Handles:
+        - Basic types: str, int, float, bool
+        - Literal types: Literal["a", "b"] → enum
+        - Union types: str | int → anyOf
+        - List types: list[str] → array
+        - Dict types: dict[str, Any] → object
+        - Callable types: Callable[..., Any] → string (module path)
+        - Annotated types: Annotated[int, Field(ge=0)] → min/max
+        - Type aliases: Resolves type aliases to their base Literal/Union types
+
+        Args
+        ----
+            type_hint: Python type annotation
+
+        Returns
+        -------
+            dict: JSON Schema definition for the type
+
+        Examples
+        --------
+        >>> from typing import Literal
+        >>> SchemaGenerator._type_to_json_schema(Literal["a", "b"])
+        {'type': 'string', 'enum': ['a', 'b']}
+        """
+        # Handle Callable types first (before checking other patterns)
+        # Callables are represented as strings (module paths) in YAML
+        if SchemaGenerator._is_callable_type(type_hint):
+            return {
+                "type": "string",
+                "description": "Module path string (e.g., 'myapp.process') or !py inline function",
+            }
+
+        if is_annotated_type(type_hint):
+            base_type, metadata = get_annotated_metadata(type_hint)
+
+            # Recursively process base type
+            schema = SchemaGenerator._type_to_json_schema(base_type)
+
+            # In Pydantic v2, Field stores metadata as FieldInfo with .metadata attribute
+            for constraint in metadata:
+                if hasattr(constraint, "metadata") and constraint.metadata:
+                    for meta_item in constraint.metadata:
+                        # Ge, Le, Gt, Lt objects
+                        if hasattr(meta_item, "ge"):
+                            schema["minimum"] = meta_item.ge
+                        if hasattr(meta_item, "le"):
+                            schema["maximum"] = meta_item.le
+                        if hasattr(meta_item, "gt"):
+                            schema["exclusiveMinimum"] = meta_item.gt
+                        if hasattr(meta_item, "lt"):
+                            schema["exclusiveMaximum"] = meta_item.lt
+                        # MinLen, MaxLen objects
+                        if hasattr(meta_item, "min_length"):
+                            schema["minLength"] = meta_item.min_length
+                        if hasattr(meta_item, "max_length"):
+                            schema["maxLength"] = meta_item.max_length
+
+                # Also check for description on the Field itself
+                if hasattr(constraint, "description") and constraint.description:
+                    schema["description"] = constraint.description
+
+            return schema
+
+        if is_literal_type(type_hint):
+            args = get_args(type_hint)
+            # Determine type from first value
+            first_val = args[0] if args else ""
+            val_type = type(first_val)
+
+            # Determine JSON Schema type from Python type
+            if val_type in SchemaGenerator.BASIC_TYPE_MAP:
+                json_type = SchemaGenerator.BASIC_TYPE_MAP[val_type]["type"]
+            else:
+                json_type = "string"
+
+            literal_schema: dict[str, Any] = {"type": json_type, "enum": list(args)}
+
+            return literal_schema
+
+        if is_union_type(type_hint):
+            args = get_args(type_hint)
+
+            # Filter out None for Optional types
+            non_none_args = [arg for arg in args if arg is not type(None)]
+
+            if len(non_none_args) == 1:
+                # Optional[T] → make nullable
+                schema = SchemaGenerator._type_to_json_schema(non_none_args[0])
+                # Allow null
+                if "type" in schema:
+                    if isinstance(schema["type"], list):
+                        schema["type"].append("null")
+                    else:
+                        schema["type"] = [schema["type"], "null"]
+                return schema
+
+            # Multiple types → anyOf
+            # But first, check if any of them are Callable and simplify
+            processed_schemas = [SchemaGenerator._type_to_json_schema(arg) for arg in non_none_args]
+
+            # Deduplicate schemas that have the same structure
+            unique_schemas = []
+            seen = set()
+            for schema in processed_schemas:
+                # Create a hashable representation
+                schema_repr = json.dumps(schema, sort_keys=True)
+                if schema_repr not in seen:
+                    seen.add(schema_repr)
+                    unique_schemas.append(schema)
+
+            if len(unique_schemas) == 1:
+                return unique_schemas[0]
+
+            return {"anyOf": unique_schemas}
+
+        if is_list_type(type_hint):
+            args = get_args(type_hint)
+            item_type = args[0] if args else Any
+
+            return {
+                "type": "array",
+                "items": SchemaGenerator._type_to_json_schema(item_type),
+            }
+
+        if is_dict_type(type_hint):
+            args = get_args(type_hint)
+            # If dict has typed values, try to extract schema
+            if len(args) >= 2:
+                value_type = args[1]
+                if value_type is not Any:
+                    return {
+                        "type": "object",
+                        "additionalProperties": SchemaGenerator._type_to_json_schema(value_type),
+                    }
+            return {"type": "object"}
+
+        if type_hint in SchemaGenerator.BASIC_TYPE_MAP:
+            return SchemaGenerator.BASIC_TYPE_MAP[type_hint].copy()
+
+        # Check if it's a type alias that we can resolve
+        # Type aliases like `Mode = Literal[...]` should be resolved
+        if hasattr(type_hint, "__value__"):
+            # Python 3.12+ type aliases have __value__
+            return SchemaGenerator._type_to_json_schema(type_hint.__value__)
+
+        # Default to string for unknown types
+        return {"type": "string"}
+
+    @staticmethod
+    def _extract_nested_structure(description: str) -> dict[str, Any] | None:
+        """Extract nested object structure from parameter description.
+
+        Parses docstring descriptions that define object structures with fields:
+        - "Each branch has: - condition: str - Expression..."
+        - "Dict with keys: field1: type - description..."
+
+        Args
+        ----
+            description: Parameter description text
+
+        Returns
+        -------
+            dict[str, Any] | None: JSON Schema properties dict if structure found, None otherwise
+
+        Examples
+        --------
+        >>> desc = "List of branches. Each has: - condition: str - The condition"
+        >>> result = SchemaGenerator._extract_nested_structure(desc)
+        >>> "condition" in result.get("properties", {}) if result else False
+        True
+        """
+        if not description:
+            return None
+
+        properties: dict[str, Any] = {}
+
+        # Pattern 1: "- field: type - description" (bullet list style)
+        # Matches lines like "- condition: str - Expression to evaluate"
+        bullet_pattern = re.compile(
+            r"-\s+(\w+):\s*(\w+)?\s*[-–—]?\s*(.*?)(?=\n\s*-|\Z)",
+            re.MULTILINE | re.DOTALL,
+        )
+        matches = bullet_pattern.findall(description)
+
+        for match in matches:
+            field_name = match[0].strip()
+            field_type = match[1].strip() if match[1] else "string"
+            field_desc = match[2].strip() if len(match) > 2 else ""
+
+            # Map common type names to JSON Schema types
+            type_map = {
+                "str": "string",
+                "string": "string",
+                "int": "integer",
+                "integer": "integer",
+                "float": "number",
+                "number": "number",
+                "bool": "boolean",
+                "boolean": "boolean",
+                "list": "array",
+                "array": "array",
+                "dict": "object",
+                "object": "object",
+            }
+
+            json_type = type_map.get(field_type.lower(), "string")
+            prop_schema: dict[str, Any] = {"type": json_type}
+            if field_desc:
+                prop_schema["description"] = field_desc
+
+            properties[field_name] = prop_schema
+
+        if properties:
+            return {"type": "object", "properties": properties}
+
+        return None
+
+    @staticmethod
+    def _extract_param_docs(func: Callable) -> dict[str, str]:
+        """Extract parameter descriptions from function docstring.
+
+        Supports multiple docstring formats:
+        - Google style (Args:)
+        - NumPy style (Parameters with --- separator)
+        - Sphinx style (:param name:)
+
+        Args
+        ----
+            func: Function to extract docs from
+
+        Returns
+        -------
+            dict[str, str]: Mapping of parameter name to description
+
+        Examples
+        --------
+        >>> def func(name: str, count: int):
+        ...     '''Function description.
+        ...
+        ...     Args:
+        ...         name: The name parameter
+        ...         count: The count parameter
+        ...     '''
+        ...     pass
+        >>> docs = SchemaGenerator._extract_param_docs(func)
+        >>> docs['name']
+        'The name parameter'
+        """
+        docstring = inspect.getdoc(func)
+        if not docstring:
+            return {}
+
+        param_docs: dict[str, str] = {}
+        lines = docstring.split("\n")
+
+        # Look for Args/Parameters section
+        in_params_section = False
+        current_param: str | None = None
+        is_numpy_style = False
+
+        for i, line in enumerate(lines):
+            line_stripped = line.strip()
+
+            # Check for NumPy-style separator (line of dashes after Parameters header)
+            # Also check if this is a short separator (4 chars like "----") which may be
+            # a formatting artifact rather than a true NumPy-style separator
+            if (
+                in_params_section
+                and line_stripped
+                and len(line_stripped) >= 3
+                and all(c in ("-", "=", "_") for c in line_stripped)
+            ):
+                # Only treat as NumPy style if the separator is long enough (>=10 chars)
+                # Short separators like "----" are often just formatting
+                if len(line_stripped) >= 10:
+                    is_numpy_style = True
+                continue
+
+            # Detect start of parameters section (case insensitive)
+            if line_stripped.lower() in (
+                "args:",
+                "arguments:",
+                "parameters:",
+                "params:",
+                "args",
+                "arguments",
+                "parameters",
+                "params",
+            ):
+                in_params_section = True
+                continue
+
+            # Exit parameters section when we hit another section header
+            if in_params_section:
+                # NumPy style: section headers are followed by separator lines
+                # Check if next line is a separator (indicating new section)
+                if i + 1 < len(lines):
+                    next_line = lines[i + 1].strip()
+                    if (
+                        line_stripped
+                        and not line.startswith((" ", "\t"))
+                        and next_line
+                        and len(next_line) >= 3
+                        and all(c in ("-", "=", "_") for c in next_line)
+                    ):
+                        break
+
+                # Google style: section headers end with :
+                if (
+                    line_stripped
+                    and line_stripped.endswith(":")
+                    and not line.startswith(" ")
+                    and line_stripped.lower() not in ("args:", "parameters:", "params:")
+                ):
+                    break
+
+                # Check for section keywords
+                if (
+                    line
+                    and not line.startswith((" ", "\t"))
+                    and line_stripped
+                    and any(
+                        keyword in line_stripped.lower()
+                        for keyword in [
+                            "returns",
+                            "raises",
+                            "yields",
+                            "examples",
+                            "notes",
+                            "see also",
+                        ]
+                    )
+                ):
+                    break
+
+            # Process parameter lines
+            if in_params_section:
+                # Sphinx style: ":param param_name: description"
+                if line_stripped.startswith(":param"):
+                    parts = line_stripped.split(":", 3)
+                    if len(parts) >= 3:
+                        param_name = parts[1].replace("param", "").strip()
+                        description = parts[2].strip()
+                        param_docs[param_name] = description
+                        current_param = param_name
+                # NumPy style: "param_name : type" on one line, description indented below
+                elif is_numpy_style and " : " in line_stripped and not line.startswith((" ", "\t")):
+                    parts = line_stripped.split(" : ", 1)
+                    param_name = parts[0].strip()
+                    # Skip type info, description comes on next indented lines
+                    if param_name and not any(
+                        keyword in param_name.lower()
+                        for keyword in ["return", "raise", "yield", "example", "note"]
+                    ):
+                        param_docs[param_name] = ""
+                        current_param = param_name
+                # NumPy style continuation: indented lines are descriptions
+                elif (
+                    is_numpy_style
+                    and current_param
+                    and line.startswith((" ", "\t"))
+                    and line_stripped
+                ):
+                    if param_docs[current_param]:
+                        param_docs[current_param] += " " + line_stripped
+                    else:
+                        param_docs[current_param] = line_stripped
+                # Google style: "param_name: description" or "param_name (type): description"
+                # Only applies when NOT in NumPy mode (to avoid conflicts)
+                elif not is_numpy_style and ":" in line_stripped:
+                    parts = line_stripped.split(":", 1)
+                    if len(parts) == 2:
+                        param_part = parts[0].strip()
+                        description = parts[1].strip()
+                        # Handle "param_name (type)" format
+                        if "(" in param_part:
+                            param_name = param_part.split("(")[0].strip()
+                        else:
+                            param_name = param_part
+                        # Make sure it's not a section header
+                        if param_name and not any(
+                            keyword in param_name.lower()
+                            for keyword in ["return", "raise", "yield", "example", "note"]
+                        ):
+                            param_docs[param_name] = description
+                            current_param = param_name
+                # Google style continuation lines (indented description for current param)
+                elif (
+                    not is_numpy_style
+                    and current_param
+                    and line.startswith((" ", "\t"))
+                    and line_stripped
+                ):
+                    if param_docs[current_param]:
+                        param_docs[current_param] += " " + line_stripped
+                    else:
+                        param_docs[current_param] = line_stripped
+
+        return param_docs
+
+    @staticmethod
+    def _format_output(schema: dict, format: str) -> dict | str:
+        """Format schema output as dict, YAML, or JSON.
+
+        Args
+        ----
+            schema: JSON Schema dict
+            format: Output format - "dict", "yaml", or "json"
+
+        Returns
+        -------
+            dict | str: Schema in requested format
+        """
+        if format == "dict":
+            return schema
+        if format == "yaml":
+            yaml_str: str = yaml.dump(schema, sort_keys=False, default_flow_style=False)
+            return yaml_str
+        if format == "json":
+            return json.dumps(schema, indent=2)
+        return schema
+
+    @staticmethod
+    def generate_example_yaml(node_type: str, schema: dict) -> str:
+        """Generate example YAML from schema.
+
+        Creates a complete YAML example with:
+        - K8s-style structure (kind, metadata, spec)
+        - Default values where available
+        - Placeholders for required fields
+        - Comments for optional fields
+
+        Args
+        ----
+            node_type: Node type name (e.g., "llm_node")
+            schema: JSON Schema dict
+
+        Returns
+        -------
+            str: Example YAML string
+
+        Examples
+        --------
+        >>> schema = {
+        ...     "properties": {
+        ...         "template": {"type": "string"},
+        ...         "model": {"type": "string", "default": "gpt-4"}
+        ...     },
+        ...     "required": ["template"]
+        ... }
+        >>> example = SchemaGenerator.generate_example_yaml("llm_node", schema)
+        >>> "kind: llm_node" in example
+        True
+        """
+        example: dict[str, Any] = {
+            "kind": node_type,
+            "metadata": {"name": f"my_{node_type}"},
+            "spec": {},
+        }
+
+        properties = schema.get("properties", {})
+        required = schema.get("required", [])
+
+        spec: dict[str, Any] = example["spec"]
+        for prop_name, prop_schema in properties.items():
+            if "default" in prop_schema:
+                # Use default value
+                spec[prop_name] = prop_schema["default"]
+            elif "examples" in prop_schema:
+                # Use first example
+                spec[prop_name] = prop_schema["examples"][0]
+            elif prop_name in required:
+                # Required field - add placeholder
+                spec[prop_name] = SchemaGenerator._placeholder_for_type(prop_schema.get("type"))
+            # Optional fields without defaults are omitted
+
+        yaml_output: str = yaml.dump(example, sort_keys=False, default_flow_style=False)
+        return yaml_output
+
+    @staticmethod
+    def _placeholder_for_type(json_type: str | list | None) -> Any:
+        """Get placeholder value for a JSON Schema type.
+
+        Args
+        ----
+            json_type: JSON Schema type string or list
+
+        Returns
+        -------
+            Any: Appropriate placeholder value
+
+        Examples
+        --------
+        >>> SchemaGenerator._placeholder_for_type("string")
+        'value'
+        >>> SchemaGenerator._placeholder_for_type("integer")
+        0
+        """
+        if isinstance(json_type, list):
+            json_type = json_type[0] if json_type else "string"
+
+        placeholders = {
+            "string": "value",
+            "integer": 0,
+            "number": 0.0,
+            "boolean": False,
+            "array": [],
+            "object": {},
+            "null": None,
+        }
+
+        return placeholders.get(json_type if isinstance(json_type, str) else "string", "value")