qtype 0.0.16__py3-none-any.whl → 0.1.1__py3-none-any.whl

qtype/semantic/generate.py

@@ -1,25 +1,42 @@
 import argparse
 import inspect
 import subprocess
+from enum import Enum
+from functools import partial
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Literal, Union, get_args, get_origin
+from typing import Annotated, Any, Literal, Union, get_args, get_origin
 
 import networkx as nx
 
+import qtype.base.types as base_types
 import qtype.dsl.model as dsl
-from qtype.dsl.validator import _is_dsl_type
+
+
+def _is_dsl_type(type_obj: Any) -> bool:
+    """Check if a type is a DSL type that should be converted to semantic."""
+    if not hasattr(type_obj, "__name__"):
+        return False
+
+    # Check if it's defined in the DSL module
+    return (
+        hasattr(type_obj, "__module__")
+        and (
+            type_obj.__module__ == dsl.__name__
+            or type_obj.__module__ == base_types.__name__
+        )
+        and not type_obj.__name__.startswith("_")
+    )
+
 
 FIELDS_TO_IGNORE = {"Application.references"}
 TYPES_TO_IGNORE = {
     "CustomType",
     "DecoderFormat",
     "Document",
-    "Flow",
     "ListType",
     "PrimitiveTypeEnum",
     "StrictBaseModel",
-    "StructuralTypeEnum",
     "TypeDefinition",
     "ToolParameter",
     "Variable",
@@ -78,7 +95,6 @@ def generate_semantic_model(args: argparse.Namespace) -> None:
             cls.__module__ == dsl.__name__
             and not name.startswith("_")
             and name not in TYPES_TO_IGNORE
-            and not name.endswith("List")
         ):
             dsl_classes.append((name, cls))
 
@@ -117,18 +133,19 @@ def generate_semantic_model(args: argparse.Namespace) -> None:
             dedent("""
             from __future__ import annotations
 
-            from typing import Any, Literal
+            from functools import partial
+            from typing import Any, Literal, Union
 
-            from pydantic import BaseModel, Field, model_validator
+            from pydantic import BaseModel, Field, RootModel
 
-            # Import enums and type aliases from DSL
+            # Import enums, mixins, and type aliases
+            from qtype.base.types import BatchableStepMixin, BatchConfig, CachedStepMixin, ConcurrencyConfig, ConcurrentStepMixin  # noqa: F401
             from qtype.dsl.model import (  # noqa: F401
                 CustomType,
                 DecoderFormat,
                 ListType,
                 PrimitiveTypeEnum,
                 StepCardinality,
-                StructuralTypeEnum,
                 ToolParameter
             )
             from qtype.dsl.model import Variable as DSLVariable  # noqa: F401
@@ -153,37 +170,18 @@ def generate_semantic_model(args: argparse.Namespace) -> None:
         # Write classes
         f.write("\n\n".join(generated))
 
-        # Write the Flow class which _could_ be generated but we want a validator to update it's carndiality
-        f.write("\n\n")
+        # Write the DocumentType
         f.write(
-            dedent('''
-            class Flow(Step):
-                """Defines a flow of steps that can be executed in sequence or parallel.
-                If input or output variables are not specified, they are inferred from
-                the first and last step, respectively.
-                """
-
-                description: str | None = Field(
-                    None, description="Optional description of the flow."
-                )
-                cardinality: StepCardinality = Field(
-                    StepCardinality.auto,
-                    description="The cardinality of the flow, inferred from its steps when set to 'auto'.",
-                )
-                mode: Literal["Complete", "Chat"] = Field("Complete")
-                steps: list[Step] = Field(..., description="List of steps or step IDs.")
-
-                @model_validator(mode="after")
-                def infer_cardinality(self) -> "Flow":
-                    if self.cardinality == StepCardinality.auto:
-                        self.cardinality = StepCardinality.one
-                        for step in self.steps:
-                            if step.cardinality == StepCardinality.many:
-                                self.cardinality = StepCardinality.many
-                                break
-                    return self
-
-        ''').lstrip()
+            dedent("""\n\n
+                DocumentType = Union[
+                    Application,
+                    AuthorizationProviderList,
+                    ModelList,
+                    ToolList,
+                    TypeList,
+                    VariableList,
+                ]
+                       """)
         )
 
     # Format the file with Ruff
@@ -201,51 +199,142 @@ def format_with_ruff(file_path: str) -> None:
     subprocess.run(["isort", file_path], check=True)
 
 
+def _get_union_args(type_annotation):
+    """Extract union args from a type, handling Annotated types."""
+    if get_origin(type_annotation) is Annotated:
+        # For Annotated[Union[...], ...], get the Union part
+        union_type = get_args(type_annotation)[0]
+        return get_args(union_type)
+    else:
+        return get_args(type_annotation)
+
+
 DSL_ONLY_UNION_TYPES = {
-    get_args(dsl.ToolType): "Tool",
-    get_args(dsl.StepType): "Step",
-    get_args(dsl.IndexType): "Index",
-    get_args(dsl.ModelType): "Model",
+    _get_union_args(dsl.ToolType): "Tool",
+    _get_union_args(dsl.StepType): "Step",
+    _get_union_args(dsl.AuthProviderType): "AuthorizationProvider",
+    _get_union_args(dsl.SecretManagerType): "SecretManager",
+    _get_union_args(dsl.SourceType): "Source",
+    _get_union_args(dsl.IndexType): "Index",
+    _get_union_args(dsl.ModelType): "Model",
 }
 
 
+def _is_dsl_only_union(args_without_str_none: tuple) -> tuple[bool, str]:
+    """
+    Check if union represents a DSL-only type pattern.
+
+    Args:
+        args_without_str_none: Union args with str and None filtered out
+
+    Returns:
+        Tuple of (is_dsl_only, semantic_type_name)
+    """
+    if args_without_str_none and args_without_str_none in DSL_ONLY_UNION_TYPES:
+        return True, DSL_ONLY_UNION_TYPES[args_without_str_none]
+    return False, ""
+
+
+def _resolve_optional_collection(args: tuple, has_none: bool) -> str | None:
+    """
+    Handle list|None -> list pattern (empty collection = None).
+
+    Args:
+        args: Union type arguments
+        has_none: Whether None is in the union
+
+    Returns:
+        Resolved type name if pattern matches, None otherwise
+    """
+    if len(args) == 2 and has_none:
+        collection_types = [
+            arg for arg in args if get_origin(arg) in {list, dict}
+        ]
+        if collection_types:
+            return dsl_to_semantic_type_name(collection_types[0])
+    return None
+
+
+def _is_id_reference_pattern(
+    args: tuple, has_str: bool, has_secret_ref: bool
+) -> bool:
+    """
+    Check if union represents an ID reference pattern (str | Type).
+
+    ID references allow DSL to use string IDs that get resolved to objects
+    in the semantic model. Exception: str | SecretReference stays as-is.
+
+    Args:
+        args: Union type arguments
+        has_str: Whether str is in the union
+        has_secret_ref: Whether SecretReference is in the union
+
+    Returns:
+        True if this is an ID reference pattern
+    """
+    return (
+        any(_is_dsl_type(arg) for arg in args)
+        and has_str
+        and not has_secret_ref
+    )
+
+
+def _strip_str_from_union(args: tuple) -> tuple:
+    """
+    Remove str component from union for ID reference pattern.
+
+    Args:
+        args: Union type arguments
+
+    Returns:
+        Args with str filtered out
+    """
+    return tuple(arg for arg in args if arg is not str)
+
+
 def _transform_union_type(args: tuple) -> str:
-    """Transform Union types, handling string ID references."""
+    """
+    Transform Union types, handling string ID references and special cases.
+
+    This function handles the semantic type generation for union types,
+    with special handling for:
+    - DSL-only types (e.g., ToolType -> Tool)
+    - ID references (str | SomeType -> SomeType)
+    - SecretReference (str | SecretReference stays as-is)
+    - Optional types (Type | None)
+
+    Args:
+        args: Tuple of types in the union
 
+    Returns:
+        String representation of the semantic type
+    """
+    # Import SecretReference for direct type comparison
+    from qtype.dsl.model import SecretReference
+
+    # Pre-compute type characteristics
     args_without_str_none = tuple(
         arg for arg in args if arg is not str and arg is not type(None)
     )
     has_none = any(arg is type(None) for arg in args)
     has_str = any(arg is str for arg in args)
+    has_secret_ref = any(arg is SecretReference for arg in args)
 
-    # First see if this is a DSL-only union type
-    # If so, just return the corresponding semantic type
-    if args_without_str_none in DSL_ONLY_UNION_TYPES:
-        if has_none:
-            # If we have a DSL type and None, we return the DSL type with None
-            return DSL_ONLY_UNION_TYPES[args_without_str_none] + " | None"
-        else:
-            # Note we don't handle the case where we have a DSL type and str,
-            # because that would indicate a reference to an ID, which we handle separately.
-            return DSL_ONLY_UNION_TYPES[args_without_str_none]
-
-    # Handle the case where we have a list | None, which in the dsl is needed, but here we will just have an empty list.
-    if len(args) == 2:
-        list_elems = [
-            arg for arg in args if get_origin(arg) in set([list, dict])
-        ]
-        if len(list_elems) > 0 and has_none:
-            # If we have a list and None, we return the list type
-            # This is to handle cases like List[SomeType] | None
-            # which in the DSL is needed, but here we will just have an empty list.
-            return dsl_to_semantic_type_name(list_elems[0])
-
-    # If the union contains a DSL type and a str, we need to drop the str
-    if any(_is_dsl_type(arg) for arg in args) and has_str:
-        # There is a DSL type and a str, which indicates something that can reference an ID.
-        # drop the str
-        args = tuple(arg for arg in args if arg is not str)
+    # Handle DSL-only union types (e.g., ToolType -> Tool)
+    is_dsl_only, dsl_semantic_name = _is_dsl_only_union(args_without_str_none)
+    if is_dsl_only:
+        return dsl_semantic_name + " | None" if has_none else dsl_semantic_name
 
+    # Handle list | None -> list (empty list is equivalent to None)
+    if resolved_collection := _resolve_optional_collection(args, has_none):
+        return resolved_collection
+
+    # Handle ID references: str | SomeType -> SomeType
+    # Exception: str | SecretReference should remain as-is
+    if _is_id_reference_pattern(args, has_str, has_secret_ref):
+        args = _strip_str_from_union(args)
+
+    # Convert remaining types to semantic type names
     return " | ".join(dsl_to_semantic_type_name(a) for a in args)
 
 
@@ -263,6 +352,27 @@ def dsl_to_semantic_type_name(field_type: Any) -> str:
     origin = get_origin(field_type)
     args = get_args(field_type)
 
+    # Handle Reference types - unwrap to get the actual type
+    # Reference[T] is a Pydantic generic model, so get_origin returns None
+    # Instead, check if Reference is in the MRO
+    if (
+        hasattr(field_type, "__mro__")
+        and base_types.Reference in field_type.__mro__
+    ):
+        # Reference[T] becomes just T in semantic model
+        # The actual type parameter is in __pydantic_generic_metadata__ for Pydantic generic models
+        if hasattr(field_type, "__pydantic_generic_metadata__"):
+            metadata = field_type.__pydantic_generic_metadata__
+            if "args" in metadata and metadata["args"]:
+                return dsl_to_semantic_type_name(metadata["args"][0])
+        return "Any"  # Fallback for untyped Reference
+
+    # Handle Annotated types - extract the underlying type
+    if origin is Annotated:
+        # For Annotated[SomeType, ...], we want to process SomeType
+        if args:
+            return dsl_to_semantic_type_name(args[0])
+
     if origin is Union or (
         hasattr(field_type, "__class__")
         and field_type.__class__.__name__ == "UnionType"
@@ -274,7 +384,13 @@ def dsl_to_semantic_type_name(field_type: Any) -> str:
         # Format literal values
         literal_values = []
         for arg in args:
-            if isinstance(arg, str):
+            if isinstance(arg, Enum):
+                # Keep the enum reference for semantic models (e.g., StepCardinality.one)
+                # not the string value
+                enum_class_name = arg.__class__.__name__
+                enum_value_name = arg.name
+                literal_values.append(f"{enum_class_name}.{enum_value_name}")
+            elif isinstance(arg, str):
                 literal_values.append(f'"{arg}"')
             else:
                 literal_values.append(str(arg))
@@ -323,16 +439,35 @@ def generate_semantic_class(class_name: str, cls: type) -> str:
     if inspect.isabstract(cls):
         inheritance += ", ABC"
 
-    # Check if this class inherits from another DSL class
+    # Collect all base classes from DSL and base_types modules
+    base_classes = []
     for base in cls.__bases__:
         if (
             hasattr(base, "__module__")
-            and base.__module__ == dsl.__name__
+            and (
+                base.__module__ == dsl.__name__
+                or base.__module__ == base_types.__name__
+            )
             and base.__name__ not in TYPES_TO_IGNORE
             and not base.__name__.startswith("_")
         ):
-            # This class inherits from another DSL class
-            semantic_base = f"{base.__name__}"
+            base_classes.append(base)
+
+    # Build inheritance string
+    if base_classes:
+        # Process DSL classes first, then mixins
+        dsl_bases = [
+            b.__name__ for b in base_classes if b.__module__ == dsl.__name__
+        ]
+        mixin_bases = [
+            b.__name__
+            for b in base_classes
+            if b.__module__ == base_types.__name__
+        ]
+
+        if dsl_bases:
+            # Inherit from the DSL class
+            semantic_base = dsl_bases[0]
             if inspect.isabstract(cls):
                 inheritance = f"ABC, {semantic_base}"
             else:
@@ -340,7 +475,15 @@ def generate_semantic_class(class_name: str, cls: type) -> str:
             if semantic_name == "Tool":
                 # Tools should inherit from Step and be immutable
                 inheritance = f"{semantic_base}, ImmutableModel"
-            break
+
+        # Add mixins to the inheritance - must come BEFORE BaseModel for correct MRO
+        if mixin_bases:
+            if inheritance == "BaseModel":
+                # Mixins must come before BaseModel
+                inheritance = f"{', '.join(mixin_bases)}, BaseModel"
+            else:
+                # If we have other bases, append mixins
+                inheritance = f"{inheritance}, {', '.join(mixin_bases)}"
 
     # Get field information from the class - only fields defined on this class, not inherited
     fields = []
@@ -354,6 +497,7 @@ def generate_semantic_class(class_name: str, cls: type) -> str:
                 field_info = cls.model_fields[field_name]
                 field_type = field_info.annotation
                 field_default = field_info.default
+                field_default_factory = field_info.default_factory
                 field_description = getattr(field_info, "description", None)
 
                 # Transform the field type
@@ -369,7 +513,11 @@ def generate_semantic_class(class_name: str, cls: type) -> str:
 
                 # Create field definition
                 field_def = create_field_definition(
-                    field_name, semantic_type, field_default, field_description
+                    field_name,
+                    semantic_type,
+                    field_default,
+                    field_default_factory,
+                    field_description,
                 )
                 fields.append(field_def)
 
@@ -391,6 +539,7 @@ def create_field_definition(
     field_name: str,
     field_type: str,
     field_default: Any,
+    field_default_factory: Any,
     field_description: str | None,
 ) -> str:
     """Create a field definition string."""
@@ -401,11 +550,29 @@ def create_field_definition(
 
     # Handle default values
     # Check for PydanticUndefined (required field)
-    from enum import Enum
-
     from pydantic_core import PydanticUndefined
 
-    if field_default is PydanticUndefined or field_default is ...:
+    # Check if there's a default_factory
+    if field_default_factory is not None:
+        # Handle default_factory - check if it's a partial
+        if isinstance(field_default_factory, partial):
+            # For partial, we need to serialize it properly
+            func_name = field_default_factory.func.__name__
+            # Get the keyword arguments from the partial
+            kwargs_str = ", ".join(
+                f"{k}={v}" if not isinstance(v, str) else f'{k}="{v}"'
+                for k, v in field_default_factory.keywords.items()
+            )
+            default_part = (
+                f"default_factory=partial({func_name}, {kwargs_str})"
+            )
+        else:
+            # Regular factory function
+            factory_name = getattr(
+                field_default_factory, "__name__", str(field_default_factory)
+            )
+            default_part = f"default_factory={factory_name}"
+    elif field_default is PydanticUndefined or field_default is ...:
         default_part = "..."
     elif field_default is None:
         default_part = "None"
@@ -430,7 +597,12 @@ def create_field_definition(
         default_part = str(field_default)
 
     # Create Field definition
-    field_parts = [default_part]
+    # If using default_factory, don't include it in field_parts list initially
+    if field_default_factory is not None:
+        field_parts = []
+    else:
+        field_parts = [default_part]
+
     if field_description:
         # Escape quotes and handle multiline descriptions
         escaped_desc = field_description.replace('"', '\\"').replace(
@@ -440,6 +612,13 @@ def create_field_definition(
     if alias_part:
         field_parts.append(alias_part.lstrip(", "))
 
-    field_def = f"Field({', '.join(field_parts)})"
+    # Handle default_factory in the Field() call
+    if field_default_factory is not None:
+        if field_parts:
+            field_def = f"Field({default_part}, {', '.join(field_parts)})"
+        else:
+            field_def = f"Field({default_part})"
+    else:
+        field_def = f"Field({', '.join(field_parts)})"
 
     return f"    {field_name}: {field_type} = {field_def}"

qtype/semantic/loader.py

@@ -0,0 +1,95 @@
+"""
+Semantic model loading and resolution.
+
+This is the main entry point for loading QType specifications.
+Coordinates the pipeline: load → parse → link → resolve → check
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from qtype.base.types import CustomTypeRegistry
+from qtype.dsl.linker import link
+from qtype.dsl.loader import load_yaml_file, load_yaml_string
+from qtype.dsl.parser import parse_document
+from qtype.semantic.checker import check
+from qtype.semantic.model import DocumentType
+from qtype.semantic.resolver import resolve
+
+
+def load(
+    source: str | Path,
+) -> tuple[DocumentType, CustomTypeRegistry]:
+    """
+    Load a QType YAML file, validate it, and return resolved semantic model.
+
+    This function coordinates the complete loading pipeline:
+    1. Load YAML (with env vars and includes)
+    2. Parse into DSL models (with custom type building)
+    3. Link references (resolve IDs to objects)
+    4. Resolve to semantic models (DSL → IR)
+    5. Check semantic rules
+
+    Args:
+        source: File path (str or Path) to load. Use load_from_string()
+                for raw YAML content.
+
+    Returns:
+        Tuple of (SemanticDocumentType, CustomTypeRegistry)
+
+    Raises:
+        YAMLLoadError: If YAML parsing fails
+        ValueError: If validation fails
+        DuplicateComponentError: If duplicate IDs found
+        ReferenceNotFoundError: If reference resolution fails
+        QTypeSemanticError: If semantic rules violated
+    """
+    # Load from file path
+    if isinstance(source, Path):
+        yaml_data = load_yaml_file(source)
+    else:
+        # Assume str is a file path
+        yaml_data = load_yaml_file(source)
+
+    dsl_doc, types = parse_document(yaml_data)
+    linked_doc = link(dsl_doc)
+    semantic_doc = resolve(linked_doc)
+    check(semantic_doc)
+    return semantic_doc, types
+
+
+def load_from_string(
+    content: str, base_path: str | Path | None = None
+) -> tuple[DocumentType, CustomTypeRegistry]:
+    """
+    Load a QType YAML from string content.
+
+    This function coordinates the complete loading pipeline:
+    1. Load YAML (with env vars and includes)
+    2. Parse into DSL models (with custom type building)
+    3. Link references (resolve IDs to objects)
+    4. Resolve to semantic models (DSL → IR)
+    5. Check semantic rules
+
+    Args:
+        content: Raw YAML content as string
+        base_path: Base path for resolving relative includes (default: cwd)
+
+    Returns:
+        Tuple of (SemanticDocumentType, CustomTypeRegistry)
+
+    Raises:
+        YAMLLoadError: If YAML parsing fails
+        ValueError: If validation fails
+        DuplicateComponentError: If duplicate IDs found
+        ReferenceNotFoundError: If reference resolution fails
+        QTypeSemanticError: If semantic rules violated
+    """
+    yaml_data = load_yaml_string(content, base_path=base_path)
+
+    dsl_doc, types = parse_document(yaml_data)
+    linked_doc = link(dsl_doc)
+    semantic_doc = resolve(linked_doc)
+    check(semantic_doc)
+    return semantic_doc, types