Tier 1 fixes for #1954 (variable composition review feedback)

docs/reference/advanced/managed-variables/templates-and-composition.md

@@ -150,16 +150,16 @@ When `safety_rules` is updated in the Logfire UI, all variables that reference `
 
 ### Composition Control Flow
 
-The `@{}@` syntax supports a small Handlebars-compatible subset for composing variables. It supports simple references, dotted field reads, and block helpers that branch or iterate over a top-level referenced variable:
+The `@{}@` syntax runs through the full Handlebars engine (just with `@{` / `}@` as the delimiter pair instead of the default `{{` / `}}`), so any expression form that works in Handlebars also works here — simple references, dotted field reads, block helpers, and helper sub-expressions:
 
 | Syntax | Description |
 |--------|-------------|
 | `@{variable_name}@` | Insert a variable's value |
 | `@{variable.field}@` | Access a nested field |
 | `@{#if variable}@...@{else}@...@{/if}@` | Conditional on whether a variable is set |
+| `@{#if user.active}@...@{/if}@` | Conditional on a dotted field |
 | `@{#each items}@...@{/each}@` | Iterate over a list variable |
-
-Block helper conditions and iterables must be top-level variable names. Use `@{#if user}@...@{user.active}@...@{/if}@` rather than `@{#if user.active}@`.
+| `@{#each items}@@{../top}@@{/each}@` | Access an outer-scope value from inside a block |
 
 ### Composition Tracking
 

logfire/_internal/main.py

@@ -2622,6 +2622,17 @@ def var(
 
         return variable
 
+    @overload
+    def template_var(
+        self,
+        name: str,
+        *,
+        default: T,
+        inputs_type: type[InputsT],
+        description: str | None = None,
+    ) -> TemplateVariable[T, InputsT]: ...
+
+    @overload
     def template_var(
         self,
         name: str,
@@ -2630,6 +2641,16 @@ def template_var(
         default: T | ResolveFunction[T],
         inputs_type: type[InputsT],
         description: str | None = None,
+    ) -> TemplateVariable[T, InputsT]: ...
+
+    def template_var(
+        self,
+        name: str,
+        *,
+        type: type[T] | None = None,
+        default: T | ResolveFunction[T],
+        inputs_type: type[InputsT],
+        description: str | None = None,
     ) -> TemplateVariable[T, InputsT]:
         """Define a managed template variable with integrated rendering.
 
@@ -2652,7 +2673,6 @@ class PromptInputs(BaseModel):
 
         prompt = logfire.template_var(
             'system_prompt',
-            type=str,
             default='Hello {{user_name}}',
             inputs_type=PromptInputs,
         )
@@ -2663,16 +2683,30 @@ class PromptInputs(BaseModel):
 
         Args:
             name: Unique identifier for the variable.
-            type: Expected type for validation and JSON schema generation.
+            type: Expected type for validation and JSON schema generation. Can be a primitive
+                type or a Pydantic model. If not provided, the type is inferred from `default`.
+                Required when `default` is a resolve function.
             default: Default value used when no remote configuration is found.
-                Can also be a callable with `targeting_key` and `attributes` parameters.
+                When `type` is not provided, the type is inferred from this value.
+                Can also be a callable with `targeting_key` and `attributes` parameters
+                (requires `type` to be set explicitly).
             inputs_type: The type (typically a Pydantic `BaseModel`) describing the expected
                 template inputs. Used for type-safe `get(inputs)` calls and JSON schema generation.
             description: Optional human-readable description of what the variable controls.
         """
         import re
 
-        from logfire.variables.variable import TemplateVariable
+        from logfire.variables.variable import TemplateVariable, is_resolve_function
+
+        if type is None:
+            if is_resolve_function(default):
+                raise TypeError(
+                    'When `default` is a resolve function (callable with targeting_key and attributes parameters), '
+                    '`type` must be provided to specify the variable value type.'
+                )
+            tp = cast(Type[T], default.__class__)  # noqa UP006
+        else:
+            tp = type
 
         if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', name):
             raise ValueError(
@@ -2692,7 +2726,7 @@ class PromptInputs(BaseModel):
 
         variable = TemplateVariable[T, InputsT](
             name,
-            type=type,
+            type=tp,
             default=default,
             inputs_type=inputs_type,
             description=description,

logfire/variables/_handlebars.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-from collections.abc import Callable
-from functools import cache
+from functools import cache, lru_cache
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from pydantic_handlebars import HandlebarsEnvironment
+    from types import ModuleType
+
+    from pydantic_handlebars import CompiledTemplate, HandlebarsEnvironment
 
 
 # The reference-syntax composition pass consumes ``@{...}@`` placeholders and
@@ -21,7 +22,7 @@ class _FallbackHandlebarsError(Exception):
 
 try:
     from pydantic_handlebars import HandlebarsError as _ImportedHandlebarsError
-except ImportError:  # pragma: no cover
+except ImportError:
     _ImportedHandlebarsError = _FallbackHandlebarsError
 
 HandlebarsError: type[Exception] = _ImportedHandlebarsError
@@ -31,16 +32,36 @@ class HandlebarsDependencyError(ImportError):
     """Raised when a Handlebars feature is used without pydantic-handlebars installed."""
 
 
-def _dependency_error() -> HandlebarsDependencyError:
-    return HandlebarsDependencyError(
-        'Handlebars template rendering requires the `pydantic-handlebars` package, '
-        'which is installed by the `logfire[variables]` extra.'
-    )
+@cache
+def _pydantic_handlebars() -> ModuleType:
+    """Return the cached `pydantic_handlebars` module, or raise a helpful error.
+
+    The import is cached so every accessor below — `get_environment`,
+    `compile_composition_template`, `extract_composition_dependencies`,
+    etc. — can grab attributes off the returned module without each one
+    repeating its own try/except dance. Failure surfaces once, here, with
+    the install hint.
+    """
+    try:
+        import pydantic_handlebars
+    except ModuleNotFoundError as exc:
+        # Only reframe the error if it's *pydantic_handlebars itself* that's
+        # missing. A future version pulling in a new transitive dep that
+        # isn't installed would raise `ModuleNotFoundError` with a different
+        # `exc.name`, and the user wants to see *that* name in the message
+        # rather than a misleading "install pydantic-handlebars" hint.
+        if exc.name != 'pydantic_handlebars':
+            raise
+        raise HandlebarsDependencyError(
+            'Handlebars template rendering requires the `pydantic-handlebars` package, '
+            'which is installed by the `logfire[variables]` extra.'
+        ) from exc
+    return pydantic_handlebars
 
 
 def ensure_handlebars_available() -> None:
     """Raise a helpful error if pydantic-handlebars is unavailable."""
-    get_environment()
+    _pydantic_handlebars()
 
 
 @cache
@@ -51,55 +72,65 @@ def get_environment() -> HandlebarsEnvironment:
     `{{...}}` runtime placeholders in the template as plain content; a
     subsequent render pass with the default delimiters consumes those.
     """
-    try:
-        from pydantic_handlebars import HandlebarsEnvironment
-    except ModuleNotFoundError as exc:  # pragma: no cover
-        if exc.name == 'pydantic_handlebars':
-            raise _dependency_error() from exc
-        raise
-    return HandlebarsEnvironment(open_delim=COMPOSITION_OPEN_DELIM, close_delim=COMPOSITION_CLOSE_DELIM)
+    return _pydantic_handlebars().HandlebarsEnvironment(
+        open_delim=COMPOSITION_OPEN_DELIM, close_delim=COMPOSITION_CLOSE_DELIM
+    )
 
 
 @cache
-def get_handlebars_renderer() -> tuple[type[str], Callable[..., str]]:
-    """Return pydantic-handlebars SafeString and module-level render function."""
-    try:
-        from pydantic_handlebars import SafeString, render
-    except ModuleNotFoundError as exc:  # pragma: no cover
-        if exc.name == 'pydantic_handlebars':
-            raise _dependency_error() from exc
-        raise
-    return SafeString, render
+def get_runtime_environment() -> HandlebarsEnvironment:
+    """Return a cached default-delimiter `HandlebarsEnvironment` for `{{...}}` rendering.
 
+    Used by `TemplateVariable.get(inputs)` to render the post-composition
+    serialized value against the provided inputs.
+    """
+    return _pydantic_handlebars().HandlebarsEnvironment()
 
-def extract_composition_dependencies(template: str) -> set[str]:
-    """Return the top-level `@{name}@` references in *template*.
 
-    Delegates to `pydantic_handlebars.extract_dependencies` configured for
-    the composition delimiters, so block helpers / dotted paths / etc. are
-    handled AST-correctly.
+@cache
+def get_safe_string_cls() -> type[str]:
+    """Return `pydantic_handlebars.SafeString`.
+
+    Context values are wrapped in it so HTML auto-escaping (off by default
+    but enableable) doesn't munge them.
     """
-    try:
-        from pydantic_handlebars import extract_dependencies
-    except ModuleNotFoundError as exc:  # pragma: no cover
-        if exc.name == 'pydantic_handlebars':
-            raise _dependency_error() from exc
-        raise
-    return extract_dependencies(template, open_delim=COMPOSITION_OPEN_DELIM, close_delim=COMPOSITION_CLOSE_DELIM)
+    return _pydantic_handlebars().SafeString
 
 
-@cache
-def _get_template_compatibility_checker() -> Callable[[list[str], dict[str, Any]], Any]:
-    """Return pydantic-handlebars schema compatibility checker, or raise a helpful error."""
-    try:
-        from pydantic_handlebars import check_template_compatibility as _check_template_compatibility
-    except ModuleNotFoundError as exc:  # pragma: no cover
-        if exc.name == 'pydantic_handlebars':
-            raise _dependency_error() from exc
-        raise
-    return _check_template_compatibility
+@lru_cache(maxsize=1024)
+def compile_composition_template(source: str) -> CompiledTemplate:
+    """Return a cached `CompiledTemplate` for *source* under composition delimiters.
+
+    Managed-variable values are typically stable across many resolutions, so
+    caching the parsed program lets `Variable._resolve` skip the parse on
+    every `get()` call. 1024 is large enough for any realistic number of
+    distinct templates in a single process while staying bounded for
+    long-running workers. Same rationale for `compile_runtime_template`.
+    """
+    return get_environment().compile(source)
+
+
+@lru_cache(maxsize=1024)
+def compile_runtime_template(source: str) -> CompiledTemplate:
+    """Return a cached `CompiledTemplate` for *source* under default `{{...}}` delimiters."""
+    return get_runtime_environment().compile(source)
+
+
+@lru_cache(maxsize=1024)
+def extract_composition_dependencies(template: str) -> set[str]:
+    """Return the top-level `@{name}@` references in *template*.
+
+    Cached because cycle / reference validation runs over the same template
+    strings multiple times per push or sync. The underlying delegation goes
+    to `pydantic_handlebars.extract_dependencies` configured for the
+    composition delimiters, so block helpers, dotted paths, and helper
+    sub-expressions are handled AST-correctly.
+    """
+    return _pydantic_handlebars().extract_dependencies(
+        template, open_delim=COMPOSITION_OPEN_DELIM, close_delim=COMPOSITION_CLOSE_DELIM
+    )
 
 
 def check_template_compatibility(templates: list[str], schema: dict[str, Any]) -> Any:
     """Run pydantic-handlebars schema compatibility checking."""
-    return _get_template_compatibility_checker()(templates, schema)
+    return _pydantic_handlebars().check_template_compatibility(templates, schema)