qtype 0.0.12__py3-none-any.whl → 0.1.7__py3-none-any.whl

qtype/dsl/linker.py

@@ -0,0 +1,384 @@
+from typing import Any, Dict, Type
+
+from pydantic import BaseModel, RootModel
+
+import qtype.base.types as base_types
+import qtype.dsl.domain_types
+import qtype.dsl.model as dsl
+
+
+class QTypeValidationError(Exception):
+    """Raised when there's an error during QType validation."""
+
+    pass
+
+
+class DuplicateComponentError(QTypeValidationError):
+    """Raised when there are duplicate components with the same ID."""
+
+    def __init__(
+        self,
+        obj_id: str,
+        found_obj: qtype.dsl.domain_types.StrictBaseModel,
+        existing_obj: qtype.dsl.domain_types.StrictBaseModel,
+    ):
+        super().__init__(
+            f"Duplicate component with ID {obj_id} found:\n"
+            + str(found_obj.model_dump_json())
+            + "\nAlready exists:\n"
+            + str(existing_obj.model_dump_json())
+        )
+
+
+class ComponentNotFoundError(QTypeValidationError):
+    """Raised when a component is not found in the DSL Application."""
+
+    def __init__(self, component_name: str):
+        super().__init__(
+            f"Component with name '{component_name}' not found in the DSL Application."
+        )
+
+
+class ReferenceNotFoundError(QTypeValidationError):
+    """Raised when a reference is not found in the lookup map."""
+
+    def __init__(
+        self,
+        reference: str,
+        type_hint: str | None = None,
+        available_refs: list[str] | None = None,
+    ):
+        if type_hint:
+            msg = (
+                f"Reference '{reference}' not found for type '{type_hint}'.\n"
+            )
+        else:
+            msg = f"Reference '{reference}' not found.\n"
+
+        # Add helpful suggestions if we have available references
+        if available_refs:
+            # Find similar names
+            similar = [
+                ref
+                for ref in available_refs
+                if reference.lower() in ref.lower()
+                or ref.lower() in reference.lower()
+            ]
+            if similar:
+                msg += f"Did you mean one of these? {', '.join(similar[:5])}"
+            elif len(available_refs) <= 10:
+                msg += f"Available references: {', '.join(available_refs)}"
+            else:
+                msg += (
+                    f"There are {len(available_refs)} available "
+                    "references. Check your spelling."
+                )
+
+        super().__init__(msg)
+
+
+def _update_map_with_unique_check(
+    current_map: Dict[str, qtype.dsl.domain_types.StrictBaseModel],
+    new_objects: list[qtype.dsl.domain_types.StrictBaseModel],
+) -> None:
+    """
+    Update a map with new objects, ensuring unique IDs.
+
+    Args:
+        current_map: The current map of objects by ID.
+        new_objects: List of new objects to add to the map.
+
+    Returns:
+        Updated map with new objects added, ensuring unique IDs.
+    """
+    for obj in new_objects:
+        if obj is None:
+            # If the object is None, we skip it.
+            continue
+        if isinstance(obj, str) or isinstance(obj, base_types.Reference):
+            # If the object is a string, we assume it is an ID and skip it.
+            # This is a special case where we do not want to add the string itself.
+            continue
+        # Note: There is no current abstraction for the `id` field, so we assume it exists.
+        obj_id = obj.id  # type: ignore[attr-defined]
+        # If the object already exists in the map, we check if it is the same object.
+        # If it is not the same object, we raise an error.
+        # This ensures that we do not have duplicate components with the same ID.
+        if obj_id in current_map and id(current_map[obj_id]) != id(obj):
+            raise DuplicateComponentError(obj.id, obj, current_map[obj_id])  # type: ignore
+        else:
+            current_map[obj_id] = obj
+
+
+def _collect_components_from_object(
+    obj: qtype.dsl.domain_types.StrictBaseModel,
+) -> list[qtype.dsl.domain_types.StrictBaseModel]:
+    """
+    Collect all components from an object that have IDs.
+    This includes the object itself and any nested components.
+
+    Args:
+        obj: The object to extract components from.
+
+    Returns:
+        List of components with IDs.
+    """
+    components = []
+
+    # Add the object itself if it has an ID
+    if hasattr(obj, "id"):
+        components.append(obj)
+
+    # For Flow, also collect embedded steps, inputs, and outputs
+    if isinstance(obj, dsl.Flow):
+        components.extend(obj.steps or [])  # type: ignore
+        components.extend(obj.variables or [])  # type: ignore
+
+    return components
+
+
+def _update_maps_with_embedded_objects(
+    lookup_map: Dict[str, qtype.dsl.domain_types.StrictBaseModel],
+    embedded_objects: list[qtype.dsl.domain_types.StrictBaseModel],
+) -> None:
+    """
+    Update lookup maps with embedded objects.
+    Embedded objects are when the user specifies the object and not just the ID.
+
+    Args:
+        lookup_maps: The current lookup maps to update.
+        embedded_objects: List of embedded objects to add to the maps.
+    """
+    for obj in embedded_objects:
+        components = _collect_components_from_object(obj)
+        _update_map_with_unique_check(lookup_map, components)
+
+
+def _build_lookup_maps(
+    document: Any,
+    lookup_map: Dict[str, qtype.dsl.domain_types.StrictBaseModel]
+    | None = None,
+) -> Dict[str, qtype.dsl.domain_types.StrictBaseModel]:
+    """
+    Build lookup map for all objects in a DSL Document.
+    This function creates a dictionary of id -> component, where each key is a
+    component id and the value is the component.
+
+    Works with any Document type (Application, Flow, *List types, etc.).
+
+    Args:
+        document: The DSL Document to build lookup maps for.
+                 Can be Application, Flow, or any RootModel list type.
+
+    Returns:
+        Dict[str, dsl.StrictBaseModel]: A dictionary of lookup maps
+
+    Throws:
+        QTypeValidationError: If there are duplicate components with the same ID.
+    """
+    if lookup_map is None:
+        lookup_map = {}
+
+    # Handle Application specially since it has multiple component lists
+    if isinstance(document, dsl.Application):
+        component_names = {
+            f
+            for f in dsl.Application.model_fields.keys()
+            if f not in {"id", "references", "description"}
+        }
+
+        for component_name in component_names:
+            if not hasattr(document, component_name):
+                raise ComponentNotFoundError(component_name)
+            components = getattr(document, component_name) or []
+            if not isinstance(components, list):
+                components = [components]  # Ensure we have a list
+            _update_map_with_unique_check(lookup_map, components)
+            _update_maps_with_embedded_objects(lookup_map, components)
+
+        # Handle references (which can contain nested Applications or other documents)
+        for ref in document.references or []:
+            ref = ref.root  # type: ignore
+            _build_lookup_maps(ref, lookup_map)
+
+        lookup_map[document.id] = document
+
+    # Handle RootModel list types (e.g., AuthorizationProviderList, IndexList, etc.)
+    elif hasattr(document, "root") and isinstance(
+        getattr(document, "root"), list
+    ):
+        root_list = getattr(document, "root")
+        _update_map_with_unique_check(lookup_map, root_list)
+        _update_maps_with_embedded_objects(lookup_map, root_list)
+
+    # Handle single component documents (e.g., Flow, Agent, etc.)
+    else:
+        components = _collect_components_from_object(document)
+        _update_map_with_unique_check(lookup_map, components)
+
+    return lookup_map
+
+
+def _resolve_reference(
+    ref: str, type_hint: Type, lookup_map: Dict[str, Any]
+) -> Any:
+    """
+    Resolve a single reference string to its object.
+
+    Args:
+        ref: The reference ID to resolve
+        type_hint: Type hint for better error messages
+        lookup_map: Map of component IDs to objects
+
+    Returns:
+        The resolved object
+
+    Raises:
+        ReferenceNotFoundError: If the reference cannot be found
+    """
+    resolved_obj = lookup_map.get(ref)
+    if resolved_obj is None:
+        available_refs = list(lookup_map.keys())
+        raise ReferenceNotFoundError(ref, str(type_hint), available_refs)
+    return resolved_obj
+
+
+def _resolve_rootmodel_references(
+    model: RootModel, lookup_map: Dict[str, Any]
+) -> None:
+    """
+    Resolve references in a RootModel (list-based documents).
+
+    Args:
+        model: RootModel instance to resolve
+        lookup_map: Map of component IDs to objects
+    """
+    root_list = model.root  # type: ignore
+    if not isinstance(root_list, list):
+        return
+
+    for i, item in enumerate(root_list):
+        match item:
+            case base_types.Reference():
+                root_list[i] = _resolve_reference(
+                    item.ref, type(item), lookup_map
+                )
+            case BaseModel():
+                _resolve_all_references(item, lookup_map)
+
+
+def _resolve_list_references(
+    field_value: list, lookup_map: Dict[str, Any]
+) -> None:
+    """
+    Resolve references within a list field.
+
+    Args:
+        field_value: List to process
+        lookup_map: Map of component IDs to objects
+    """
+    for i, item in enumerate(field_value):
+        match item:
+            case base_types.Reference():
+                field_value[i] = _resolve_reference(
+                    item.ref, type(item), lookup_map
+                )
+            case BaseModel():
+                _resolve_all_references(item, lookup_map)
+
+
+def _resolve_dict_references(
+    field_value: dict, lookup_map: Dict[str, Any]
+) -> None:
+    """
+    Resolve references within a dict field.
+
+    Args:
+        field_value: Dict to process
+        lookup_map: Map of component IDs to objects
+    """
+    for k, v in field_value.items():
+        match v:
+            case base_types.Reference():
+                field_value[k] = _resolve_reference(v.ref, type(v), lookup_map)
+            case BaseModel():
+                _resolve_all_references(v, lookup_map)
+
+
+def _resolve_all_references(
+    model: BaseModel,
+    lookup_map: Dict[str, Any],
+) -> None:
+    """
+    Walk a Pydantic model tree and resolve all Reference objects.
+
+    Args:
+        model: The model to process
+        lookup_map: Map of component IDs to objects
+    """
+    # Check if this is a RootModel (list-based document like ModelList, ToolList, etc.)
+    if isinstance(model, RootModel):
+        _resolve_rootmodel_references(model, lookup_map)
+        return
+
+    # For regular BaseModel types, iterate over fields
+    for field_name, field_value in model.__iter__():
+        match field_value:
+            case base_types.Reference():
+                setattr(
+                    model,
+                    field_name,
+                    _resolve_reference(
+                        field_value.ref, type(field_value), lookup_map
+                    ),
+                )
+            case BaseModel():
+                _resolve_all_references(field_value, lookup_map)
+            case list() if len(field_value) > 0:
+                _resolve_list_references(field_value, lookup_map)
+            case dict():
+                _resolve_dict_references(field_value, lookup_map)
+
+
+def link(document: dsl.DocumentType) -> dsl.DocumentType:
+    """
+    Links (resolves) all ID references in a DSL Document to their actual objects.
+
+    Works with any DocumentType:
+    - Application: Full application with all components
+    - Flow: Individual flow definition
+    - Agent: Individual agent definition
+    - AuthorizationProviderList: List of authorization providers
+    - IndexList: List of indexes
+    - ModelList: List of models
+    - ToolList: List of tools
+    - TypeList: List of custom types
+    - VariableList: List of variables
+
+    IMPORTANT: The returned object breaks the type safety of the original.
+    All Reference[T] fields will be replaced with actual T objects, which
+    violates the original type signatures. This is intentional for the
+    linking phase before transformation to semantic IR.
+
+    Args:
+        document: Any valid DSL DocumentType (one of the 9 possible document structures).
+
+    Returns:
+        The same document with all internal references resolved to actual objects.
+
+    Raises:
+        DuplicateComponentError: If there are duplicate components with the same ID.
+        ReferenceNotFoundError: If a reference cannot be resolved.
+        ComponentNotFoundError: If an expected component is missing.
+    """
+
+    # First, make a lookup map of all objects in the document.
+    # This ensures that all object ids are unique.
+    lookup_map = _build_lookup_maps(document)
+
+    # Now we resolve all ID references in the document.
+    # All DocumentType variants are BaseModel instances (including RootModel-based *List types)
+    if isinstance(document, BaseModel):
+        _resolve_all_references(document, lookup_map)
+
+    return document  # type: ignore[return-value]

qtype/dsl/loader.py

@@ -0,0 +1,315 @@
+"""
+YAML loading with environment variable and file inclusion support.
+
+This module provides two explicit functions for loading YAML:
+- load_yaml_file(path): Load YAML from a file path or URI
+- load_yaml_string(content, base_path): Load YAML from a string
+
+Both support:
+- Environment variable substitution (${VAR} or ${VAR:default})
+- File inclusion (!include and !include_raw)
+- Multiple URI schemes via fsspec (local, http, s3, etc.)
+
+Example:
+    # Load from file
+    data = load_yaml_file("config.yaml")
+    data = load_yaml_file("s3://bucket/config.yaml")
+
+    # Load from string
+    yaml_content = "name: test\\nvalue: ${ENV_VAR}"
+    data = load_yaml_string(yaml_content)
+
+    # Load string with includes (requires base_path)
+    data = load_yaml_string(yaml_content, base_path="/path/to/configs")
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+from typing import Any
+
+import fsspec
+import yaml
+from dotenv import load_dotenv
+
+
+class YAMLLoadError(Exception):
+    """Error during YAML loading or parsing."""
+
+    def __init__(
+        self,
+        message: str,
+        line: int | None = None,
+        column: int | None = None,
+        source: str | None = None,
+        original_error: Exception | None = None,
+    ) -> None:
+        self.message = message
+        self.line = line
+        self.column = column
+        self.source = source
+        self.original_error = original_error
+        super().__init__(self._format_message())
+
+    def _format_message(self) -> str:
+        """Format error message with location information."""
+        parts = []
+        if self.source:
+            parts.append(f"in {self.source}")
+        if self.line is not None:
+            location = f"line {self.line + 1}"
+            if self.column is not None:
+                location += f", column {self.column + 1}"
+            parts.append(location)
+
+        if parts:
+            return f"{self.message} ({', '.join(parts)})"
+        return self.message
+
+
+class YAMLLoader(yaml.SafeLoader):
+    """YAML loader with env var substitution and file inclusion."""
+
+    def __init__(self, stream: Any, base_path: str | None = None) -> None:
+        super().__init__(stream)
+        self.base_path = base_path or str(Path.cwd())
+
+
+def _substitute_env_vars(value: str) -> str:
+    """
+    Substitute environment variables in a string.
+
+    Supports ${VAR_NAME} or ${VAR_NAME:default} syntax.
+
+    Args:
+        value: String containing environment variable references
+
+    Returns:
+        String with environment variables substituted
+
+    Raises:
+        ValueError: If required environment variable is not found
+    """
+    pattern = r"\$\{([^}:]+)(?::([^}]*))?\}"
+
+    def replace_env_var(match: re.Match[str]) -> str:
+        var_name = match.group(1)
+        default_value = match.group(2)
+
+        env_value = os.getenv(var_name)
+
+        if env_value is not None:
+            return env_value
+        elif default_value is not None:
+            return default_value
+        else:
+            raise ValueError(
+                f"Environment variable '{var_name}' is required but not set"
+            )
+
+    return re.sub(pattern, replace_env_var, value)
+
+
+def _resolve_path(base_path: str, target_path: str) -> str:
+    """
+    Resolve a target path relative to base path.
+
+    Uses fsspec's URL joining logic which handles both local paths and URIs.
+
+    Args:
+        base_path: Base path or URI
+        target_path: Target path to resolve (relative or absolute)
+
+    Returns:
+        Resolved absolute path or URI
+    """
+    # If target is already absolute (has scheme or starts with /), use as-is
+    from urllib.parse import urljoin, urlparse
+
+    parsed = urlparse(target_path)
+    if parsed.scheme or target_path.startswith("/"):
+        return target_path
+
+    # Check if base is URL-like
+    base_parsed = urlparse(base_path)
+    if base_parsed.scheme:
+        # URL-based resolution
+        return urljoin(base_path, target_path)
+    else:
+        # Local file resolution
+        base_path_obj = Path(base_path)
+        if not base_path_obj.is_dir():
+            base_path_obj = base_path_obj.parent
+        return str(base_path_obj / target_path)
+
+
+def _env_var_constructor(loader: YAMLLoader, node: yaml.ScalarNode) -> str:
+    """Constructor for environment variable substitution."""
+    value = loader.construct_scalar(node)
+    return _substitute_env_vars(value)
+
+
+def _include_constructor(loader: YAMLLoader, node: yaml.ScalarNode) -> Any:
+    """Constructor for !include tag to load external YAML files."""
+    file_path = loader.construct_scalar(node)
+    resolved_path = _resolve_path(loader.base_path, file_path)
+
+    try:
+        with fsspec.open(resolved_path, "r", encoding="utf-8") as f:
+            content = f.read()  # type: ignore[misc]
+            # Create a partial function to pass base_path to YAMLLoader
+            from functools import partial
+
+            loader_class = partial(YAMLLoader, base_path=resolved_path)
+            return yaml.load(content, loader_class)  # type: ignore[arg-type]
+    except (FileNotFoundError, IOError, OSError) as e:
+        raise FileNotFoundError(
+            f"Failed to load included file '{resolved_path}': {e}"
+        ) from e
+
+
+def _include_raw_constructor(loader: YAMLLoader, node: yaml.ScalarNode) -> str:
+    """Constructor for !include_raw tag to load external text files."""
+    file_path = loader.construct_scalar(node)
+    resolved_path = _resolve_path(loader.base_path, file_path)
+
+    try:
+        with fsspec.open(resolved_path, "r", encoding="utf-8") as f:
+            return f.read()  # type: ignore[no-any-return]
+    except (FileNotFoundError, IOError, OSError) as e:
+        raise FileNotFoundError(
+            f"Failed to load included file '{resolved_path}': {e}"
+        ) from e
+
+
+# Register constructors
+YAMLLoader.add_constructor("tag:yaml.org,2002:str", _env_var_constructor)
+YAMLLoader.add_constructor("!include", _include_constructor)
+YAMLLoader.add_constructor("!include_raw", _include_raw_constructor)
+
+
+def load_yaml_file(path: str | Path) -> dict[str, Any]:
+    """
+    Load YAML from a file path or URI.
+
+    Supports multiple URI schemes via fsspec (local files, http, s3, etc.).
+    Automatically loads .env files from the source directory.
+
+    Args:
+        path: File path or URI to load
+
+    Returns:
+        Parsed YAML as dictionary
+
+    Raises:
+        YAMLLoadError: If YAML parsing fails
+        FileNotFoundError: If file doesn't exist
+        ValueError: If required environment variable is missing
+    """
+    source_str = str(path)
+
+    # Load .env file if it exists in the source directory
+    try:
+        from urllib.parse import urlparse
+
+        parsed = urlparse(source_str)
+        if parsed.scheme in ["file", ""]:
+            # Local file - load .env from same directory
+            source_path = Path(parsed.path if parsed.path else source_str)
+            if source_path.is_file():
+                env_dir = source_path.parent
+                env_file = env_dir / ".env"
+                if env_file.exists():
+                    load_dotenv(env_file)
+    except Exception:
+        pass
+
+    # Also try cwd
+    load_dotenv()
+
+    # Load file content
+    try:
+        with fsspec.open(source_str, "r", encoding="utf-8") as f:
+            content = f.read()  # type: ignore[misc]
+    except FileNotFoundError as e:
+        raise FileNotFoundError(f"File not found: {source_str}") from e
+
+    return _parse_yaml(content, base_path=source_str, source_name=source_str)
+
+
+def load_yaml_string(
+    content: str, base_path: str | Path | None = None
+) -> dict[str, Any]:
+    """
+    Load YAML from a string.
+
+    Args:
+        content: Raw YAML content as string
+        base_path: Base path for resolving relative includes (default: cwd)
+
+    Returns:
+        Parsed YAML as dictionary
+
+    Raises:
+        YAMLLoadError: If YAML parsing fails
+        ValueError: If required environment variable is missing
+    """
+    load_dotenv()
+
+    base = str(base_path) if base_path else str(Path.cwd())
+    return _parse_yaml(content, base_path=base, source_name="<string>")
+
+
+def _parse_yaml(
+    content: str, base_path: str, source_name: str
+) -> dict[str, Any]:
+    """
+    Parse YAML content with environment variable substitution and includes.
+
+    Args:
+        content: YAML content to parse
+        base_path: Base path for resolving relative includes
+        source_name: Source name for error messages
+
+    Returns:
+        Parsed YAML as dictionary
+
+    Raises:
+        YAMLLoadError: If YAML parsing fails
+    """
+    try:
+        from functools import partial
+
+        loader_class = partial(YAMLLoader, base_path=base_path)
+        result = yaml.load(content, loader_class)  # type: ignore[arg-type]
+        return result  # type: ignore[no-any-return]
+    except yaml.YAMLError as e:
+        # Extract line/column information if available
+        line = None
+        column = None
+
+        if hasattr(e, "problem_mark") and e.problem_mark:  # type: ignore[attr-defined]
+            line = e.problem_mark.line  # type: ignore[attr-defined]
+            column = e.problem_mark.column  # type: ignore[attr-defined]
+
+        # Format error message
+        error_msg = str(e)
+        if hasattr(e, "problem"):
+            error_msg = e.problem or error_msg  # type: ignore[attr-defined]
+
+        raise YAMLLoadError(
+            message=f"YAML parsing error: {error_msg}",
+            line=line,
+            column=column,
+            source=source_name,
+            original_error=e,
+        ) from e
+    except ValueError as e:
+        # Environment variable errors
+        raise YAMLLoadError(
+            message=str(e),
+            source=source_name,
+            original_error=e,
+        ) from e