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

ostruct/cli/cli.py

@@ -43,6 +43,8 @@ from openai import (
 )
 from openai_structured.client import (
     async_openai_structured_stream,
+    get_context_window_limit,
+    get_default_token_limit,
     supports_structured_output,
 )
 from openai_structured.errors import (
@@ -441,48 +443,6 @@ def estimate_tokens_for_chat(
         return num_tokens
 
 
-def get_default_token_limit(model: str) -> int:
-    """Get the default token limit for a given model.
-
-    Note: These limits are based on current OpenAI model specifications as of 2024 and may
-    need to be updated if OpenAI changes the models' capabilities.
-
-    Args:
-        model: The model name (e.g., 'gpt-4o', 'o1-mini', 'o3-mini')
-
-    Returns:
-        The default token limit for the model
-    """
-    if "o1-" in model:
-        return 100_000  # o1-mini supports up to 100K output tokens
-    elif "gpt-4o" in model:
-        return 16_384  # gpt-4o supports up to 16K output tokens
-    elif "o3-" in model:
-        return 16_384  # o3-mini supports up to 16K output tokens
-    else:
-        return 4_096  # default fallback
-
-
-def get_context_window_limit(model: str) -> int:
-    """Get the total context window limit for a given model.
-
-    Note: These limits are based on current OpenAI model specifications as of 2024 and may
-    need to be updated if OpenAI changes the models' capabilities.
-
-    Args:
-        model: The model name (e.g., 'gpt-4o', 'o1-mini', 'o3-mini')
-
-    Returns:
-        The context window limit for the model
-    """
-    if "o1-" in model:
-        return 200_000  # o1-mini supports 200K total context window
-    elif "gpt-4o" in model or "o3-" in model:
-        return 128_000  # gpt-4o and o3-mini support 128K context window
-    else:
-        return 8_192  # default fallback
-
-
 def validate_token_limits(
     model: str, total_tokens: int, max_token_limit: Optional[int] = None
 ) -> None:
@@ -771,7 +731,9 @@ def _validate_path_mapping_internal(
             raise
 
         if security_manager:
-            if not security_manager.is_allowed_file(str(resolved_path)):
+            try:
+                security_manager.validate_path(str(resolved_path))
+            except PathSecurityError:
                 raise PathSecurityError.from_expanded_paths(
                     original_path=str(path),
                     expanded_path=str(resolved_path),
@@ -1192,40 +1154,13 @@ def validate_security_manager(
     if base_dir is None:
         base_dir = os.getcwd()
 
-    # Default to empty list if allowed_dirs is None
-    if allowed_dirs is None:
-        allowed_dirs = []
-
-    # Add base directory if it exists
-    try:
-        base_dir_path = Path(base_dir).resolve()
-        if not base_dir_path.exists():
-            raise DirectoryNotFoundError(
-                f"Base directory not found: {base_dir}"
-            )
-        if not base_dir_path.is_dir():
-            raise DirectoryNotFoundError(
-                f"Base directory is not a directory: {base_dir}"
-            )
-        all_allowed_dirs = [str(base_dir_path)]
-    except OSError as e:
-        raise DirectoryNotFoundError(f"Invalid base directory: {e}")
+    # Create security manager with base directory
+    security_manager = SecurityManager(base_dir)
 
     # Add explicitly allowed directories
-    for dir_path in allowed_dirs:
-        try:
-            resolved_path = Path(dir_path).resolve()
-            if not resolved_path.exists():
-                raise DirectoryNotFoundError(
-                    f"Directory not found: {dir_path}"
-                )
-            if not resolved_path.is_dir():
-                raise DirectoryNotFoundError(
-                    f"Path is not a directory: {dir_path}"
-                )
-            all_allowed_dirs.append(str(resolved_path))
-        except OSError as e:
-            raise DirectoryNotFoundError(f"Invalid directory path: {e}")
+    if allowed_dirs:
+        for dir_path in allowed_dirs:
+            security_manager.add_allowed_directory(dir_path)
 
     # Add directories from file if specified
     if allowed_dir_file:
@@ -1234,28 +1169,13 @@ def validate_security_manager(
                 for line in f:
                     line = line.strip()
                     if line and not line.startswith("#"):
-                        try:
-                            resolved_path = Path(line).resolve()
-                            if not resolved_path.exists():
-                                raise DirectoryNotFoundError(
-                                    f"Directory not found: {line}"
-                                )
-                            if not resolved_path.is_dir():
-                                raise DirectoryNotFoundError(
-                                    f"Path is not a directory: {line}"
-                                )
-                            all_allowed_dirs.append(str(resolved_path))
-                        except OSError as e:
-                            raise DirectoryNotFoundError(
-                                f"Invalid directory path in {allowed_dir_file}: {e}"
-                            )
+                        security_manager.add_allowed_directory(line)
         except OSError as e:
             raise DirectoryNotFoundError(
                 f"Failed to read allowed directories file: {e}"
             )
 
-    # Create security manager with all allowed directories
-    return SecurityManager(base_dir=base_dir, allowed_dirs=all_allowed_dirs)
+    return security_manager
 
 
 def parse_var(var_str: str) -> Tuple[str, str]:
@@ -1416,29 +1336,78 @@ async def stream_structured_output(
     It handles the core streaming logic and resource cleanup.
     """
     try:
-        async for chunk in async_openai_structured_stream(
-            client=client,
-            model=model,
-            output_schema=output_schema,
-            system_prompt=system_prompt,
-            user_prompt=user_prompt,
-            **kwargs,
-        ):
-            if not chunk:
-                continue
-
-            # Process and output the chunk
-            dumped = chunk.model_dump(mode="json")
-            json_str = json.dumps(dumped, indent=2)
-
-            if output_file:
-                with open(output_file, "a", encoding="utf-8") as f:
-                    f.write(json_str)
-                    f.write("\n")
-                    f.flush()  # Ensure immediate flush to file
+        # Base models that don't support streaming
+        non_streaming_models = {"o1", "o3"}
+
+        # Check if model supports streaming
+        # o3-mini and o3-mini-high support streaming, base o3 does not
+        use_streaming = model not in non_streaming_models and (
+            not model.startswith("o3") or model.startswith("o3-mini")
+        )
+
+        # All o1 and o3 models (base and variants) have fixed settings
+        stream_kwargs = {}
+        if not (model.startswith("o1") or model.startswith("o3")):
+            stream_kwargs = kwargs
+
+        if use_streaming:
+            async for chunk in async_openai_structured_stream(
+                client=client,
+                model=model,
+                output_schema=output_schema,
+                system_prompt=system_prompt,
+                user_prompt=user_prompt,
+                **stream_kwargs,
+            ):
+                if not chunk:
+                    continue
+
+                # Process and output the chunk
+                dumped = chunk.model_dump(mode="json")
+                json_str = json.dumps(dumped, indent=2)
+
+                if output_file:
+                    with open(output_file, "a", encoding="utf-8") as f:
+                        f.write(json_str)
+                        f.write("\n")
+                        f.flush()  # Ensure immediate flush to file
+                else:
+                    # Print directly to stdout with immediate flush
+                    print(json_str, flush=True)
+        else:
+            # For non-streaming models, use regular completion
+            response = await client.chat.completions.create(
+                model=model,
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": user_prompt},
+                ],
+                stream=False,
+                **stream_kwargs,
+            )
+
+            # Process the single response
+            content = response.choices[0].message.content
+            if content:
+                try:
+                    # Parse and validate against schema
+                    result = output_schema.model_validate_json(content)
+                    json_str = json.dumps(
+                        result.model_dump(mode="json"), indent=2
+                    )
+
+                    if output_file:
+                        with open(output_file, "w", encoding="utf-8") as f:
+                            f.write(json_str)
+                            f.write("\n")
+                    else:
+                        print(json_str, flush=True)
+                except ValidationError as e:
+                    raise InvalidResponseFormatError(
+                        f"Response validation failed: {e}"
+                    )
             else:
-                # Print directly to stdout with immediate flush
-                print(json_str, flush=True)
+                raise EmptyResponseError("Model returned empty response")
 
     except (
         StreamInterruptedError,
@@ -1657,8 +1626,7 @@ def create_cli() -> click.Command:
             logger.exception("Unexpected error")
             raise CLIError(str(e), context={"error_type": type(e).__name__})
 
-    # The decorated function is a Command, but mypy can't detect this
-    return cast(click.Command, cli)
+    return cli
 
 
 def main() -> None:
@@ -1671,8 +1639,6 @@ def main() -> None:
 __all__ = [
     "ExitCode",
     "estimate_tokens_for_chat",
-    "get_context_window_limit",
-    "get_default_token_limit",
     "parse_json_var",
     "create_dynamic_model",
     "validate_path_mapping",

ostruct/cli/click_options.py

@@ -1,17 +1,14 @@
 """Click command and options for the CLI.
 
 This module contains all Click-related code separated from the main CLI logic.
-We isolate this code here and disable mypy type checking for the entire module
-because Click's decorator-based API is not easily type-checkable, leading to
-many type: ignore comments in the main code.
+We isolate this code here and provide proper type annotations for Click's
+decorator-based API.
 """
 
-# mypy: ignore-errors
-# ^ This tells mypy to ignore type checking for this entire file
-
-from typing import Any, Callable
+from typing import Any, Callable, TypeVar, Union, cast
 
 import click
+from click import Command
 
 from ostruct import __version__
 from ostruct.cli.errors import (  # noqa: F401 - Used in error handling
@@ -19,6 +16,9 @@ from ostruct.cli.errors import (  # noqa: F401 - Used in error handling
     TaskTemplateVariableError,
 )
 
+F = TypeVar("F", bound=Callable[..., Any])
+DecoratedCommand = Union[Command, Callable[..., Any]]
+
 
 def validate_task_params(
     ctx: click.Context, param: click.Parameter, value: Any
@@ -162,87 +162,96 @@ def model_options(f: Callable) -> Callable:
     return f
 
 
-def create_click_command() -> Callable:
-    """Create the Click command with all options."""
+def create_click_command() -> Callable[[F], Command]:
+    """Create the Click command with all options.
+
+    Returns:
+        A decorator function that adds all CLI options to the command.
+    """
+
+    def decorator(f: F) -> Command:
+        # Start with the base command
+        cmd: DecoratedCommand = click.command(
+            help="Make structured OpenAI API calls."
+        )(f)
 
-    def decorator(f: Callable) -> Callable:
-        f = click.command(help="Make structured OpenAI API calls.")(f)
-        f = click.option(
+        # Add all options
+        cmd = click.option(
             "--task",
             help="Task template string",
             type=str,
             callback=validate_task_params,
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--task-file",
             help="Task template file path",
             type=str,
             callback=validate_task_params,
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--system-prompt",
             help="System prompt string",
             type=str,
             callback=validate_system_prompt_params,
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--system-prompt-file",
             help="System prompt file path",
             type=str,
             callback=validate_system_prompt_params,
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--schema-file",
             required=True,
             help="JSON schema file for response validation",
             type=str,
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--ignore-task-sysprompt",
             is_flag=True,
             help="Ignore system prompt from task template YAML frontmatter",
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--timeout",
             type=float,
             default=60.0,
             help="API timeout in seconds",
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--output-file", help="Write JSON output to file", type=str
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--dry-run",
             is_flag=True,
             help="Simulate API call without making request",
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--no-progress", is_flag=True, help="Disable progress indicators"
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--progress-level",
             type=click.Choice(["none", "basic", "detailed"]),
             default="basic",
             help="Progress reporting level",
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--api-key", help="OpenAI API key (overrides env var)", type=str
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--verbose",
             is_flag=True,
             help="Enable verbose output and detailed logging",
-        )(f)
-        f = click.option(
+        )(cmd)
+        cmd = click.option(
             "--debug-openai-stream",
             is_flag=True,
             help="Enable low-level debug output for OpenAI streaming",
-        )(f)
-        f = debug_options(f)
-        f = file_options(f)
-        f = variable_options(f)
-        f = model_options(f)
-        f = click.version_option(version=__version__)(f)
-        return f
+        )(cmd)
+        cmd = debug_options(cmd)
+        cmd = file_options(cmd)
+        cmd = variable_options(cmd)
+        cmd = model_options(cmd)
+        cmd = click.version_option(version=__version__)(cmd)
+        return cast(Command, cmd)
 
     return decorator