ostruct-cli 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

ostruct/cli/security/security_manager.py

@@ -15,6 +15,8 @@ from contextlib import contextmanager
 from pathlib import Path
 from typing import Generator, List, Optional, Union
 
+from ostruct.cli.errors import OstructFileNotFoundError
+
 from .allowed_checker import is_path_in_allowed_dirs
 from .case_manager import CaseManager
 from .errors import (
@@ -200,24 +202,32 @@ class SecurityManager:
             PathSecurityError: If the path fails security validation
             FileNotFoundError: If the file doesn't exist (only checked after security validation)
         """
+        logger.debug("Validating path: %s", path)
+
         # First normalize the path
         norm_path = normalize_path(path)
+        logger.debug("Normalized path: %s", norm_path)
 
         # Handle symlinks first - delegate to symlink_resolver
         if norm_path.is_symlink():
+            logger.debug("Path is a symlink, resolving: %s", norm_path)
             try:
-                return _resolve_symlink(
+                resolved = _resolve_symlink(
                     norm_path,
                     self._max_symlink_depth,
                     self._allowed_dirs,
                 )
+                logger.debug("Resolved symlink to: %s", resolved)
+                return resolved
             except RuntimeError as e:
                 if "Symlink loop" in str(e):
+                    logger.error("Symlink loop detected: %s", path)
                     raise PathSecurityError(
                         "Symlink security violation: loop detected",
                         path=str(path),
                         context={"reason": SecurityErrorReasons.SYMLINK_LOOP},
                     ) from e
+                logger.error("Failed to resolve symlink: %s - %s", path, e)
                 raise PathSecurityError(
                     f"Symlink security violation: failed to resolve symlink - {e}",
                     path=str(path),
@@ -225,10 +235,13 @@ class SecurityManager:
                 ) from e
 
         # For non-symlinks, just check if the normalized path is allowed
+        logger.debug("Checking if path is allowed: %s", norm_path)
         if not self.is_path_allowed(norm_path):
             logger.error(
-                "Security violation: Path %s is outside allowed directories",
+                "Security violation: Path %s is outside allowed directories (base_dir=%s, allowed_dirs=%s)",
                 path,
+                self._base_dir,
+                self._allowed_dirs,
             )
             raise PathSecurityError(
                 (
@@ -244,12 +257,12 @@ class SecurityManager:
             )
 
         # Only check existence after security validation passes
+        logger.debug("Checking if path exists: %s", norm_path)
         if not norm_path.exists():
             logger.debug("Path allowed but not found: %s", norm_path)
-            raise FileNotFoundError(
-                f"File not found: {os.path.basename(str(path))}"
-            )
+            raise OstructFileNotFoundError(str(path))
 
+        logger.debug("Path validation successful: %s", norm_path)
         return norm_path
 
     def resolve_path(self, path: Union[str, Path]) -> Path:
@@ -275,7 +288,7 @@ class SecurityManager:
             if self._allow_temp_paths and self.is_temp_path(norm_path):
                 logger.debug("Allowing temp path: %s", norm_path)
                 if not norm_path.exists():
-                    raise FileNotFoundError(f"File not found: {path}")
+                    raise OstructFileNotFoundError(f"File not found: {path}")
                 return norm_path
 
             # Handle symlinks with security checks
@@ -296,7 +309,7 @@ class SecurityManager:
                     elif reason == SecurityErrorReasons.SYMLINK_BROKEN:
                         msg = f"Broken symlink: {e.context['source']} -> {e.context['target']}"
                         logger.debug(msg)
-                        raise FileNotFoundError(msg) from e
+                        raise OstructFileNotFoundError(msg) from e
                     # Any other security errors propagate unchanged
                     raise
 
@@ -318,12 +331,12 @@ class SecurityManager:
 
             # Only check existence after security validation
             if not norm_path.exists():
-                raise FileNotFoundError(f"File not found: {path}")
+                raise OstructFileNotFoundError(f"File not found: {path}")
 
             return norm_path
 
         except OSError as e:
-            if isinstance(e, FileNotFoundError):
+            if isinstance(e, OstructFileNotFoundError):
                 raise
             logger.error("Error resolving path: %s - %s", path, e)
             raise PathSecurityError(

ostruct/cli/serialization.py

@@ -0,0 +1,25 @@
+"""Serialization utilities for CLI logging."""
+
+import json
+from typing import Any, Dict
+
+
+class LogSerializer:
+    """Utility class for serializing log data."""
+
+    @staticmethod
+    def serialize_log_extra(extra: Dict[str, Any]) -> str:
+        """Serialize extra log data to a formatted string.
+
+        Args:
+            extra: Dictionary of extra log data
+
+        Returns:
+            Formatted string representation of the extra data
+        """
+        try:
+            # Try to serialize with nice formatting
+            return json.dumps(extra, indent=2, default=str)
+        except Exception:
+            # Fall back to basic string representation if JSON fails
+            return str(extra)

ostruct/cli/template_filters.py

@@ -10,7 +10,7 @@ from collections import Counter
 from typing import Any, Dict, List, Optional, Sequence, TypeVar, Union
 
 import tiktoken
-from jinja2 import Environment
+from jinja2 import Environment, pass_context
 from pygments import highlight
 from pygments.formatters import HtmlFormatter, NullFormatter, TerminalFormatter
 from pygments.lexers import TextLexer, get_lexer_by_name, guess_lexer
@@ -178,10 +178,12 @@ def format_error(e: Exception) -> str:
     return f"{type(e).__name__}: {str(e)}"
 
 
-def estimate_tokens(text: str) -> int:
+@pass_context
+def estimate_tokens(context: Any, text: str) -> int:
     """Estimate number of tokens in text."""
     try:
-        encoding = tiktoken.encoding_for_model("gpt-4")
+        # Use o200k_base encoding for token estimation
+        encoding = tiktoken.get_encoding("o200k_base")
         return len(encoding.encode(str(text)))
     except Exception as e:
         logger.warning(f"Failed to estimate tokens: {e}")

ostruct/cli/template_rendering.py

@@ -61,8 +61,9 @@ from typing import Any, Dict, List, Optional, Union
 import jinja2
 from jinja2 import Environment
 
-from .errors import TemplateValidationError
+from .errors import TaskTemplateVariableError, TemplateValidationError
 from .file_utils import FileInfo
+from .progress import ProgressContext
 from .template_env import create_jinja_env
 from .template_schema import DotDict, StdinProxy
 
@@ -92,23 +93,23 @@ TemplateContextValue = Union[
 def render_template(
     template_str: str,
     context: Dict[str, Any],
-    jinja_env: Optional[Environment] = None,
-    progress_enabled: bool = True,
+    env: Optional[Environment] = None,
+    progress: Optional[ProgressContext] = None,
 ) -> str:
-    """Render a task template with the given context.
+    """Render a template with the given context.
 
     Args:
-        template_str: Task template string or path to task template file
-        context: Task template variables
-        jinja_env: Optional Jinja2 environment to use
-        progress_enabled: Whether to show progress indicators
+        template_str: Template string to render
+        context: Context dictionary for template variables
+        env: Optional Jinja2 environment to use
+        progress: Optional progress bar to update
 
     Returns:
-        Rendered task template string
+        str: The rendered template string
 
     Raises:
-        TemplateValidationError: If task template cannot be loaded or rendered. The original error
-                  will be chained using `from` for proper error context.
+        TaskTemplateVariableError: If template variables are undefined
+        TemplateValidationError: If template rendering fails for other reasons
     """
     from .progress import (  # Import here to avoid circular dependency
         ProgressContext,
@@ -116,16 +117,14 @@ def render_template(
 
     with ProgressContext(
         description="Rendering task template",
-        level="basic" if progress_enabled else "none",
+        level="basic" if progress else "none",
     ) as progress:
         try:
             if progress:
                 progress.update(1)  # Update progress for setup
 
-            if jinja_env is None:
-                jinja_env = create_jinja_env(
-                    loader=jinja2.FileSystemLoader(".")
-                )
+            if env is None:
+                env = create_jinja_env(loader=jinja2.FileSystemLoader("."))
 
             logger.debug("=== Raw Input ===")
             logger.debug(
@@ -154,7 +153,7 @@ def render_template(
             # Wrap JSON variables in DotDict and handle special cases
             wrapped_context: Dict[str, TemplateContextValue] = {}
             for key, value in context.items():
-                if isinstance(value, dict):
+                if isinstance(value, dict) and not isinstance(value, DotDict):
                     wrapped_context[key] = DotDict(value)
                 else:
                     wrapped_context[key] = value
@@ -188,7 +187,7 @@ def render_template(
                         f"Task template file not found: {template_str}"
                     )
                 try:
-                    template = jinja_env.get_template(template_str)
+                    template = env.get_template(template_str)
                 except jinja2.TemplateNotFound as e:
                     raise TemplateValidationError(
                         f"Task template file not found: {e.name}"
@@ -199,10 +198,10 @@ def render_template(
                     template_str,
                 )
                 try:
-                    template = jinja_env.from_string(template_str)
+                    template = env.from_string(template_str)
 
                     # Add debug log for loop rendering
-                    def debug_file_render(f: FileInfo) -> str:
+                    def debug_file_render(f: FileInfo) -> Any:
                         logger.info("Rendering file: %s", f.path)
                         return ""
 
@@ -292,6 +291,10 @@ def render_template(
                             "  %s: %s (%r)", key, type(value).__name__, value
                         )
                 result = template.render(**wrapped_context)
+                if not isinstance(result, str):
+                    raise TemplateValidationError(
+                        f"Template rendered to non-string type: {type(result)}"
+                    )
                 logger.info(
                     "Template render result (first 100 chars): %r",
                     result[:100],
@@ -302,7 +305,17 @@ def render_template(
                 )
                 if progress:
                     progress.update(1)
-                return result  # type: ignore[no-any-return]
+                return result
+            except jinja2.UndefinedError as e:
+                # Extract variable name from error message
+                var_name = str(e).split("'")[1]
+                error_msg = (
+                    f"Missing required template variable: {var_name}\n"
+                    f"Available variables: {', '.join(sorted(context.keys()))}\n"
+                    "To fix this, please provide the variable using:\n"
+                    f"  -V {var_name}='value'"
+                )
+                raise TaskTemplateVariableError(error_msg) from e
             except (jinja2.TemplateError, Exception) as e:
                 logger.error("Template rendering failed: %s", str(e))
                 raise TemplateValidationError(
@@ -336,4 +349,15 @@ def render_template_file(
     """
     with open(template_path, "r", encoding="utf-8") as f:
         template_str = f.read()
-    return render_template(template_str, context, jinja_env, progress_enabled)
+
+    # Create a progress context if enabled
+    progress = (
+        ProgressContext(
+            description="Rendering template file",
+            level="basic" if progress_enabled else "none",
+        )
+        if progress_enabled
+        else None
+    )
+
+    return render_template(template_str, context, jinja_env, progress)

ostruct/cli/template_utils.py

@@ -9,10 +9,11 @@ It re-exports the public APIs from specialized modules:
 - template_io: File I/O operations and metadata extraction
 """
 
+import logging
 from typing import Any, Dict, List, Optional, Set
 
 import jsonschema
-from jinja2 import Environment, meta
+from jinja2 import Environment, TemplateSyntaxError, meta
 from jinja2.nodes import Node
 
 from .errors import (
@@ -35,6 +36,8 @@ from .template_schema import (
 )
 from .template_validation import SafeUndefined, validate_template_placeholders
 
+logger = logging.getLogger(__name__)
+
 
 # Custom error classes
 class TemplateMetadataError(TaskTemplateError):
@@ -118,8 +121,13 @@ def find_all_template_variables(
     if env is None:
         env = Environment(undefined=SafeUndefined)
 
-    # Parse template
-    ast = env.parse(template)
+    try:
+        ast = env.parse(template)
+    except TemplateSyntaxError as e:
+        logger.error("Failed to parse template: %s", str(e))
+        return set()  # Return empty set on parse error
+
+    variables: Set[str] = set()
 
     # Find all variables in this template
     variables = meta.find_undeclared_variables(ast)
@@ -252,7 +260,7 @@ def find_all_template_variables(
                 visit_nodes(node.else_)
 
     visit_nodes(ast.body)
-    return variables  # type: ignore[no-any-return]
+    return variables
 
 
 __all__ = [

ostruct/cli/template_validation.py

@@ -67,7 +67,7 @@ from jinja2 import meta
 from jinja2.nodes import For, Name, Node
 
 from . import template_filters
-from .errors import TemplateValidationError
+from .errors import TaskTemplateVariableError, TemplateValidationError
 from .template_env import create_jinja_env
 from .template_schema import (
     DictProxy,
@@ -160,7 +160,8 @@ def validate_template_placeholders(
         template_context: Optional context to validate against
 
     Raises:
-        TemplateValidationError: If any placeholders are invalid
+        TaskTemplateVariableError: If any variables are undefined
+        TemplateValidationError: If template validation fails for other reasons
     """
     logger = logging.getLogger(__name__)
 
@@ -309,10 +310,15 @@ def validate_template_placeholders(
         }
 
         if missing:
-            raise TemplateValidationError(
-                f"Task template uses undefined variables: {', '.join(sorted(missing))}. "
-                f"Available variables: {', '.join(sorted(available_vars))}"
+            # Create a more user-friendly error message
+            error_msg = (
+                f"Missing required template variable(s): {', '.join(sorted(missing))}\n"
+                f"Available variables: {', '.join(sorted(available_vars))}\n"
+                "To fix this, please provide the missing variable(s) using:\n"
             )
+            for var in sorted(missing):
+                error_msg += f"  -V {var}='value'\n"
+            raise TaskTemplateVariableError(error_msg)
 
         logger.debug(
             "Before create_validation_context - available_vars type: %s, value: %s",
@@ -336,8 +342,17 @@ def validate_template_placeholders(
         try:
             env.from_string(template).render(validation_context)
         except jinja2.UndefinedError as e:
-            # Convert Jinja2 undefined errors to our own ValueError
-            raise TemplateValidationError(str(e))
+            # Convert Jinja2 undefined errors to TaskTemplateVariableError with helpful message
+            var_name = str(e).split("'")[
+                1
+            ]  # Extract variable name from error message
+            error_msg = (
+                f"Missing required template variable: {var_name}\n"
+                f"Available variables: {', '.join(sorted(available_vars))}\n"
+                "To fix this, please provide the variable using:\n"
+                f"  -V {var_name}='value'"
+            )
+            raise TaskTemplateVariableError(error_msg) from e
         except ValueError as e:
             # Convert validation errors from template_schema to TemplateValidationError
             raise TemplateValidationError(str(e))
@@ -348,10 +363,13 @@ def validate_template_placeholders(
             raise
 
     except jinja2.TemplateSyntaxError as e:
-        # Convert Jinja2 syntax errors to our own ValueError
+        # Convert Jinja2 syntax errors to TemplateValidationError
         raise TemplateValidationError(
             f"Invalid task template syntax: {str(e)}"
         )
+    except (TaskTemplateVariableError, TemplateValidationError):
+        # Re-raise these without wrapping
+        raise
     except Exception as e:
         logger.error("Unexpected error during template validation: %s", str(e))
         raise

ostruct/cli/token_utils.py

@@ -0,0 +1,43 @@
+"""Token estimation utilities."""
+
+from typing import Any, Dict, List, Union
+
+import tiktoken
+
+
+def estimate_tokens_with_encoding(
+    messages: Union[str, Dict[str, str], List[Dict[str, str]]],
+    model: str,
+    encoder: Any = None,
+) -> int:
+    """Estimate the number of tokens in a chat completion.
+
+    Args:
+        messages: Message content - can be string, single message dict, or list of messages
+        model: Model name
+        encoder: Optional tiktoken encoder for testing
+
+    Returns:
+        int: Estimated token count
+    """
+    if encoder is None:
+        # Use o200k_base for gpt-4o and o1 models
+        if model.startswith(("gpt-4o", "o1", "o3")):
+            encoder = tiktoken.get_encoding("o200k_base")
+        else:
+            encoder = tiktoken.get_encoding("cl100k_base")
+
+    if isinstance(messages, str):
+        return len(encoder.encode(messages))
+    elif isinstance(messages, dict):
+        return len(encoder.encode(str(messages.get("content", ""))))
+    else:
+        num_tokens = 0
+        for message in messages:
+            num_tokens += 4  # message overhead
+            for key, value in message.items():
+                num_tokens += len(encoder.encode(str(value)))
+                if key == "name":
+                    num_tokens -= 1  # role is omitted
+        num_tokens += 2  # reply priming
+        return num_tokens

ostruct/cli/validators.py

@@ -0,0 +1,109 @@
+"""Validators for CLI options and arguments."""
+
+import json
+from pathlib import Path
+from typing import Any, List, Optional, Tuple, Union
+
+import click
+
+from .errors import InvalidJSONError, VariableNameError
+
+
+def validate_name_path_pair(
+    ctx: click.Context,
+    param: click.Parameter,
+    value: List[Tuple[str, Union[str, Path]]],
+) -> List[Tuple[str, Union[str, Path]]]:
+    """Validate name/path pairs for files and directories.
+
+    Args:
+        ctx: Click context
+        param: Click parameter
+        value: List of (name, path) tuples
+
+    Returns:
+        List of validated (name, Path) tuples
+
+    Raises:
+        click.BadParameter: If validation fails
+    """
+    if not value:
+        return value
+
+    result: List[Tuple[str, Union[str, Path]]] = []
+    for name, path in value:
+        if not name.isidentifier():
+            raise click.BadParameter(f"Invalid variable name: {name}")
+        result.append((name, Path(path)))
+    return result
+
+
+def validate_variable(
+    ctx: click.Context, param: click.Parameter, value: Optional[List[str]]
+) -> Optional[List[Tuple[str, str]]]:
+    """Validate name=value format for simple variables.
+
+    Args:
+        ctx: Click context
+        param: Click parameter
+        value: List of "name=value" strings
+
+    Returns:
+        List of validated (name, value) tuples
+
+    Raises:
+        click.BadParameter: If validation fails
+    """
+    if not value:
+        return None
+
+    result = []
+    for var in value:
+        if "=" not in var:
+            raise click.BadParameter(
+                f"Variable must be in format name=value: {var}"
+            )
+        name, val = var.split("=", 1)
+        if not name.isidentifier():
+            raise click.BadParameter(f"Invalid variable name: {name}")
+        result.append((name, val))
+    return result
+
+
+def validate_json_variable(
+    ctx: click.Context, param: click.Parameter, value: Optional[List[str]]
+) -> Optional[List[Tuple[str, Any]]]:
+    """Validate JSON variable format.
+
+    Args:
+        ctx: Click context
+        param: Click parameter
+        value: List of "name=json_string" values
+
+    Returns:
+        List of validated (name, parsed_json) tuples
+
+    Raises:
+        click.BadParameter: If validation fails
+    """
+    if not value:
+        return None
+
+    result = []
+    for var in value:
+        if "=" not in var:
+            raise InvalidJSONError(
+                f'JSON variable must be in format name=\'{"json":"value"}\': {var}'
+            )
+        name, json_str = var.split("=", 1)
+        if not name.isidentifier():
+            raise VariableNameError(f"Invalid variable name: {name}")
+        try:
+            json_value = json.loads(json_str)
+            result.append((name, json_value))
+        except json.JSONDecodeError as e:
+            raise InvalidJSONError(
+                f"Invalid JSON value for variable {name!r}: {json_str!r}",
+                context={"variable_name": name},
+            ) from e
+    return result