nighthawk-python 0.1.0__py3-none-any.whl

nighthawk/runtime/scoping.py

@@ -0,0 +1,288 @@
+from __future__ import annotations
+
+import uuid
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry.trace import get_tracer_provider
+
+from ..configuration import StepExecutorConfiguration, StepExecutorConfigurationPatch
+from ..errors import NighthawkError
+from ..tools.registry import tool_scope
+
+if TYPE_CHECKING:
+    from .step_executor import AgentStepExecutor, StepExecutor
+
+
+@dataclass(frozen=True)
+class ExecutionContext:
+    """Immutable snapshot of the current execution context.
+
+    Attributes:
+        run_id: Unique identifier for the run.
+        scope_id: Unique identifier for the current scope.
+    """
+
+    run_id: str
+    scope_id: str
+
+
+RUN_ID = "run.id"
+SCOPE_ID = "scope.id"
+STEP_ID = "step.id"
+TOOL_CALL_ID = "tool_call.id"
+
+
+_tracer = get_tracer_provider().get_tracer("nighthawk")
+
+
+@contextmanager
+def span(span_name: str, /, **attributes: Any) -> Iterator[None]:
+    with _tracer.start_as_current_span(span_name, attributes=attributes):
+        yield
+
+
+def _generate_id() -> str:
+    return uuid.uuid4().hex
+
+
+_step_executor_var: ContextVar[StepExecutor | None] = ContextVar(
+    "nighthawk_step_executor",
+    default=None,
+)
+
+_execution_context_var: ContextVar[ExecutionContext | None] = ContextVar(
+    "nighthawk_execution_context",
+    default=None,
+)
+
+
+_system_prompt_suffix_fragments_var: ContextVar[tuple[str, ...]] = ContextVar(
+    "nighthawk_system_prompt_suffix_fragments",
+    default=(),
+)
+
+_user_prompt_suffix_fragments_var: ContextVar[tuple[str, ...]] = ContextVar(
+    "nighthawk_user_prompt_suffix_fragments",
+    default=(),
+)
+
+
+def get_step_executor() -> StepExecutor:
+    """Return the active step executor.
+
+    Raises:
+        NighthawkError: If no step executor is set (i.e. called outside a run context).
+    """
+    step_executor = _step_executor_var.get()
+    if step_executor is None:
+        raise NighthawkError("StepExecutor is not set")
+    return step_executor
+
+
+def get_execution_context() -> ExecutionContext:
+    """Return the active execution context.
+
+    Raises:
+        NighthawkError: If no execution context is set (i.e. called outside a run context).
+    """
+    execution_context = _execution_context_var.get()
+    if execution_context is None:
+        raise NighthawkError("ExecutionContext is not set")
+    return execution_context
+
+
+def get_system_prompt_suffix_fragments() -> tuple[str, ...]:
+    return _system_prompt_suffix_fragments_var.get()
+
+
+def get_user_prompt_suffix_fragments() -> tuple[str, ...]:
+    return _user_prompt_suffix_fragments_var.get()
+
+
+def _resolve_agent_step_executor(step_executor: StepExecutor) -> AgentStepExecutor:
+    from .step_executor import AgentStepExecutor
+
+    if not isinstance(step_executor, AgentStepExecutor):
+        raise NighthawkError("StepExecutor configuration updates require current step_executor to be AgentStepExecutor")
+    return step_executor
+
+
+def _replace_step_executor_with_configuration(
+    step_executor: StepExecutor,
+    *,
+    configuration: StepExecutorConfiguration,
+) -> StepExecutor:
+    from .step_executor import AgentStepExecutor
+
+    current_step_executor = _resolve_agent_step_executor(step_executor)
+
+    if current_step_executor.agent_is_managed:
+        return AgentStepExecutor.from_configuration(configuration=configuration)
+
+    if current_step_executor.agent is None:
+        raise NighthawkError("AgentStepExecutor.agent is not initialized")
+    return AgentStepExecutor.from_agent(
+        agent=current_step_executor.agent,
+        configuration=configuration,
+    )
+
+
+@contextmanager
+def run(
+    step_executor: StepExecutor,
+    *,
+    run_id: str | None = None,
+) -> Iterator[None]:
+    """Start an execution run with the given step executor.
+
+    Establishes a run-scoped context that makes the step executor
+    available to all Natural blocks executed within this scope.
+
+    Args:
+        step_executor: The step executor to use for Natural block execution.
+        run_id: Optional identifier for the run. If not provided, a UUID is
+            generated automatically.
+
+    Yields:
+        None
+
+    Example:
+        ```python
+        executor = AgentStepExecutor.from_configuration(
+            configuration=StepExecutorConfiguration(model="openai:gpt-4o"),
+        )
+        with nighthawk.run(executor):
+            result = my_natural_function()
+        ```
+    """
+    execution_context = ExecutionContext(
+        run_id=run_id or _generate_id(),
+        scope_id=_generate_id(),
+    )
+
+    with tool_scope():
+        step_executor_token = _step_executor_var.set(step_executor)
+        execution_context_token = _execution_context_var.set(execution_context)
+        system_fragments_token = _system_prompt_suffix_fragments_var.set(())
+        user_fragments_token = _user_prompt_suffix_fragments_var.set(())
+        try:
+            with span(
+                "nighthawk.run",
+                **{
+                    RUN_ID: execution_context.run_id,
+                    SCOPE_ID: execution_context.scope_id,
+                },
+            ):
+                yield
+        finally:
+            _user_prompt_suffix_fragments_var.reset(user_fragments_token)
+            _system_prompt_suffix_fragments_var.reset(system_fragments_token)
+            _execution_context_var.reset(execution_context_token)
+            _step_executor_var.reset(step_executor_token)
+
+
+@contextmanager
+def scope(
+    *,
+    step_executor_configuration: StepExecutorConfiguration | None = None,
+    step_executor_configuration_patch: StepExecutorConfigurationPatch | None = None,
+    step_executor: StepExecutor | None = None,
+    system_prompt_suffix_fragment: str | None = None,
+    user_prompt_suffix_fragment: str | None = None,
+) -> Iterator[StepExecutor]:
+    """Open a nested scope that can override the step executor or its configuration.
+
+    Must be called inside an active run context. Creates a new scope_id while
+    inheriting the run_id from the parent context.
+
+    Args:
+        step_executor_configuration: Full replacement configuration for the step
+            executor.
+        step_executor_configuration_patch: Partial override applied on top of the
+            current configuration.
+        step_executor: Replacement step executor for this scope.
+        system_prompt_suffix_fragment: Additional text appended to the system prompt.
+        user_prompt_suffix_fragment: Additional text appended to the user prompt.
+
+    Yields:
+        The step executor active within this scope.
+
+    Example:
+        ```python
+        with nighthawk.run(executor):
+            with nighthawk.scope(
+                step_executor_configuration_patch=StepExecutorConfigurationPatch(
+                    model="openai:gpt-4o-mini",
+                ),
+            ) as scoped_executor:
+                result = my_natural_function()
+        ```
+    """
+    current_step_executor = get_step_executor()
+    current_execution_context = get_execution_context()
+
+    next_step_executor = current_step_executor
+
+    if step_executor is not None:
+        next_step_executor = step_executor
+
+    has_configuration_update = any(
+        value is not None
+        for value in (
+            step_executor_configuration,
+            step_executor_configuration_patch,
+        )
+    )
+
+    if has_configuration_update:
+        current_agent_step_executor = _resolve_agent_step_executor(next_step_executor)
+        next_configuration = current_agent_step_executor.configuration
+
+        if step_executor_configuration is not None:
+            next_configuration = step_executor_configuration
+
+        if step_executor_configuration_patch is not None:
+            next_configuration = step_executor_configuration_patch.apply_to(next_configuration)
+
+        next_step_executor = _replace_step_executor_with_configuration(
+            next_step_executor,
+            configuration=next_configuration,
+        )
+
+    next_execution_context = replace(
+        current_execution_context,
+        scope_id=_generate_id(),
+    )
+
+    next_system_fragments = _system_prompt_suffix_fragments_var.get()
+    next_user_fragments = _user_prompt_suffix_fragments_var.get()
+
+    if system_prompt_suffix_fragment is not None:
+        next_system_fragments = (*next_system_fragments, system_prompt_suffix_fragment)
+
+    if user_prompt_suffix_fragment is not None:
+        next_user_fragments = (*next_user_fragments, user_prompt_suffix_fragment)
+
+    with tool_scope():
+        step_executor_token = _step_executor_var.set(next_step_executor)
+        execution_context_token = _execution_context_var.set(next_execution_context)
+        system_fragments_token = _system_prompt_suffix_fragments_var.set(next_system_fragments)
+        user_fragments_token = _user_prompt_suffix_fragments_var.set(next_user_fragments)
+        try:
+            with span(
+                "nighthawk.scope",
+                **{
+                    RUN_ID: next_execution_context.run_id,
+                    SCOPE_ID: next_execution_context.scope_id,
+                },
+            ):
+                yield next_step_executor
+        finally:
+            _user_prompt_suffix_fragments_var.reset(user_fragments_token)
+            _system_prompt_suffix_fragments_var.reset(system_fragments_token)
+            _execution_context_var.reset(execution_context_token)
+            _step_executor_var.reset(step_executor_token)

nighthawk/runtime/step_context.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field, replace
+from types import CellType
+
+from ..errors import NighthawkError
+from ..json_renderer import JsonRendererStyle
+
+_MISSING: object = object()
+"""Sentinel indicating a name could not be resolved. Distinct from None."""
+
+
+@dataclass(frozen=True)
+class ToolResultRenderingPolicy:
+    tokenizer_encoding_name: str
+    tool_result_max_tokens: int
+    json_renderer_style: JsonRendererStyle
+
+
+DEFAULT_TOOL_RESULT_RENDERING_POLICY = ToolResultRenderingPolicy(
+    tokenizer_encoding_name="o200k_base",
+    tool_result_max_tokens=2_000,
+    json_renderer_style="strict",
+)
+
+
+@dataclass
+class StepContext:
+    """Mutable, per-step execution context passed to tools and executors.
+
+    ``step_globals`` and ``step_locals`` are mutable dicts. All mutations to
+    ``step_locals`` MUST go through :meth:`record_assignment` (for top-level
+    name bindings) or through the dotted-path assignment in
+    ``tools.assignment`` (which bumps ``step_locals_revision`` directly).
+    Direct dict writes bypass revision tracking and ``assigned_binding_names``
+    bookkeeping, which will cause incorrect commit behavior at Natural block
+    boundaries.
+    """
+
+    step_id: str
+
+    step_globals: dict[str, object]
+    step_locals: dict[str, object]
+
+    binding_commit_targets: set[str]
+    read_binding_names: frozenset[str]
+
+    # Ordinary user-provided binding (for example a global named "memory") may exist in step_locals.
+
+    binding_name_to_type: dict[str, object] = field(default_factory=dict)
+    assigned_binding_names: set[str] = field(default_factory=set)
+    step_locals_revision: int = 0
+    tool_result_rendering_policy: ToolResultRenderingPolicy | None = None
+
+    def record_assignment(self, name: str, value: object) -> None:
+        """Record an assignment to a step local variable.
+
+        Updates step_locals, marks the name as assigned, and bumps the revision.
+        """
+        self.step_locals[name] = value
+        self.assigned_binding_names.add(name)
+        self.step_locals_revision += 1
+
+
+_step_context_stack_var: ContextVar[tuple[StepContext, ...]] = ContextVar(
+    "nighthawk_step_context_stack",
+    default=(),
+)
+
+
+@dataclass(frozen=True)
+class PythonLookupState:
+    python_name_scope_stack: tuple[dict[str, object], ...] = ()
+    python_cell_scope_stack: tuple[dict[str, CellType], ...] = ()
+
+
+_python_lookup_state_var: ContextVar[PythonLookupState] = ContextVar(
+    "nighthawk_python_lookup_state",
+    default=PythonLookupState(),  # noqa: B039
+)
+
+
+@contextmanager
+def step_context_scope(step_context: StepContext) -> Iterator[None]:
+    current_stack = _step_context_stack_var.get()
+    token = _step_context_stack_var.set((*current_stack, step_context))
+    try:
+        yield
+    finally:
+        _step_context_stack_var.reset(token)
+
+
+@contextmanager
+def python_name_scope(name_to_value: dict[str, object]) -> Iterator[None]:
+    current_lookup_state = _python_lookup_state_var.get()
+    next_lookup_state = replace(
+        current_lookup_state,
+        python_name_scope_stack=(*current_lookup_state.python_name_scope_stack, dict(name_to_value)),
+    )
+    token = _python_lookup_state_var.set(next_lookup_state)
+    try:
+        yield
+    finally:
+        _python_lookup_state_var.reset(token)
+
+
+@contextmanager
+def python_cell_scope(name_to_cell: dict[str, CellType]) -> Iterator[None]:
+    current_lookup_state = _python_lookup_state_var.get()
+    next_lookup_state = replace(
+        current_lookup_state,
+        python_cell_scope_stack=(*current_lookup_state.python_cell_scope_stack, dict(name_to_cell)),
+    )
+    token = _python_lookup_state_var.set(next_lookup_state)
+    try:
+        yield
+    finally:
+        _python_lookup_state_var.reset(token)
+
+
+def get_step_context_stack() -> tuple[StepContext, ...]:
+    return _step_context_stack_var.get()
+
+
+def get_python_name_scope_stack() -> tuple[dict[str, object], ...]:
+    return _python_lookup_state_var.get().python_name_scope_stack
+
+
+def get_python_cell_scope_stack() -> tuple[dict[str, CellType], ...]:
+    return _python_lookup_state_var.get().python_cell_scope_stack
+
+
+def get_current_step_context() -> StepContext:
+    """Return the innermost active step context.
+
+    Raises:
+        NighthawkError: If no step context is set (i.e. called outside step execution).
+    """
+    stack = _step_context_stack_var.get()
+    if not stack:
+        raise NighthawkError("StepContext is not set")
+    return stack[-1]
+
+
+def resolve_name_in_step_context(step_context: StepContext, name: str) -> object:
+    """Resolve a name from step locals, step globals, or builtins.
+
+    Returns the resolved value, or ``_MISSING`` if the name cannot be found.
+    Callers must compare the result with ``value is _MISSING`` to detect
+    absent names (since the resolved value itself may be ``None``).
+    """
+    if name in step_context.step_locals:
+        return step_context.step_locals[name]
+
+    if name in step_context.step_globals:
+        return step_context.step_globals[name]
+
+    python_builtins = step_context.step_globals.get("__builtins__", __builtins__)
+
+    if isinstance(python_builtins, dict):
+        if name in python_builtins:
+            return python_builtins[name]
+        return _MISSING
+
+    if hasattr(python_builtins, name):
+        return getattr(python_builtins, name)
+
+    return _MISSING

nighthawk/runtime/step_contract.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+from typing import Annotated, Literal, get_args
+
+from pydantic import BaseModel, ConfigDict, Field
+
+type StepKind = Literal["pass", "return", "break", "continue", "raise"]
+
+STEP_KINDS: tuple[StepKind, ...] = get_args(StepKind.__value__)
+
+
+_REFERENCE_PATH_PATTERN = r"^(?!__)[A-Za-z_][A-Za-z0-9_]*(?:\.(?!__)[A-Za-z_][A-Za-z0-9_]*)*$"
+
+
+class PassStepOutcome(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    kind: Literal["pass"]
+
+
+class ReturnStepOutcome(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    kind: Literal["return"]
+    return_reference_path: str
+
+
+class BreakStepOutcome(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    kind: Literal["break"]
+
+
+class ContinueStepOutcome(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    kind: Literal["continue"]
+
+
+class RaiseStepOutcome(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    kind: Literal["raise"]
+    raise_message: str
+    raise_error_type: str | None = None
+
+
+# This union is used for host-side parsing after unwrapping the StepFinalResult envelope.
+
+type StepOutcome = Annotated[
+    PassStepOutcome | ReturnStepOutcome | BreakStepOutcome | ContinueStepOutcome | RaiseStepOutcome,
+    Field(discriminator="kind"),
+]
+
+
+# NOTE: StepFinalResult is an envelope that wraps StepOutcome for LLM structured output.
+# OpenAI / Codex structured outputs require the root JSON schema to be a single object (no anyOf at root level).
+# StepOutcome is a discriminated union whose JSON schema produces anyOf at root, which is rejected by these providers.
+# By placing the union inside a ``result`` field, the anyOf moves to a nested level where it is accepted.
+# Additionally, each variant uses ``additionalProperties: false`` so that kind-specific properties are enforced (e.g. ``kind: "return"`` cannot include ``raise_message``).
+
+
+class StepFinalResult(BaseModel):
+    """Envelope that wraps StepOutcome for LLM structured output."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    result: Annotated[
+        PassStepOutcome | ReturnStepOutcome | BreakStepOutcome | ContinueStepOutcome | RaiseStepOutcome,
+        Field(discriminator="kind"),
+    ]
+
+
+def _build_pass_variant_schema() -> dict[str, object]:
+    return {
+        "type": "object",
+        "properties": {
+            "kind": {"type": "string", "const": "pass"},
+        },
+        "required": ["kind"],
+        "additionalProperties": False,
+    }
+
+
+def _build_return_variant_schema() -> dict[str, object]:
+    return {
+        "type": "object",
+        "properties": {
+            "kind": {"type": "string", "const": "return"},
+            "return_reference_path": {"type": "string", "pattern": _REFERENCE_PATH_PATTERN},
+        },
+        "required": ["kind", "return_reference_path"],
+        "additionalProperties": False,
+    }
+
+
+def _build_break_variant_schema() -> dict[str, object]:
+    return {
+        "type": "object",
+        "properties": {
+            "kind": {"type": "string", "const": "break"},
+        },
+        "required": ["kind"],
+        "additionalProperties": False,
+    }
+
+
+def _build_continue_variant_schema() -> dict[str, object]:
+    return {
+        "type": "object",
+        "properties": {
+            "kind": {"type": "string", "const": "continue"},
+        },
+        "required": ["kind"],
+        "additionalProperties": False,
+    }
+
+
+def _build_raise_variant_schema(
+    *,
+    raise_error_type_binding_names: tuple[str, ...],
+) -> dict[str, object]:
+    properties: dict[str, object] = {
+        "kind": {"type": "string", "const": "raise"},
+        "raise_message": {"type": "string"},
+    }
+    required: list[str] = ["kind", "raise_message"]
+
+    if raise_error_type_binding_names:
+        properties["raise_error_type"] = {
+            "type": "string",
+            "enum": list(raise_error_type_binding_names),
+        }
+
+    return {
+        "type": "object",
+        "properties": properties,
+        "required": required,
+        "additionalProperties": False,
+    }
+
+
+def _build_variant_schema(
+    kind: StepKind,
+    *,
+    raise_error_type_binding_names: tuple[str, ...],
+) -> dict[str, object]:
+    match kind:
+        case "pass":
+            return _build_pass_variant_schema()
+        case "return":
+            return _build_return_variant_schema()
+        case "break":
+            return _build_break_variant_schema()
+        case "continue":
+            return _build_continue_variant_schema()
+        case "raise":
+            return _build_raise_variant_schema(raise_error_type_binding_names=raise_error_type_binding_names)
+
+
+def build_step_json_schema(
+    *,
+    allowed_kinds: tuple[StepKind, ...],
+    raise_error_type_binding_names: tuple[str, ...],
+) -> dict[str, object]:
+    if not allowed_kinds:
+        raise ValueError("allowed_kinds must not be empty")
+
+    variants: list[dict[str, object]] = [
+        _build_variant_schema(kind, raise_error_type_binding_names=raise_error_type_binding_names) for kind in allowed_kinds
+    ]
+
+    if len(variants) == 1:
+        result_schema: dict[str, object] = variants[0]
+    else:
+        result_schema = {"anyOf": variants}
+
+    schema: dict[str, object] = {
+        "type": "object",
+        "title": "StepFinalResult",
+        "properties": {"result": result_schema},
+        "required": ["result"],
+        "additionalProperties": False,
+    }
+
+    return schema
+
+
+def build_step_system_prompt_suffix_fragment(
+    *,
+    allowed_kinds: tuple[StepKind, ...],
+    raise_error_type_binding_names: tuple[str, ...],
+) -> str:
+    if not allowed_kinds:
+        raise ValueError("allowed_kinds must not be empty")
+
+    allowed_kinds_text = ", ".join(f"`{outcome_kind}`" for outcome_kind in allowed_kinds)
+
+    sections: list[str] = []
+    sections.append(
+        f"StepFinalResult — output exactly one JSON object with a `result` field. Inside `result`, `kind` must be one of: {allowed_kinds_text}.\n"
+    )
+
+    if "pass" in allowed_kinds:
+        sections.append('\nDefault: {"result": {"kind": "pass"}}\nChoose pass after completing the work. Most blocks end with pass.\n')
+
+    alternatives: list[str] = []
+
+    if "return" in allowed_kinds:
+        alternatives.append(
+            '- return: immediately return from the Python function.\n  return_reference_path (required): name in step locals holding the value.\n  Example: after nh_assign("x", "42"), output {"result": {"kind": "return", "return_reference_path": "x"}}.\n'
+        )
+
+    if "break" in allowed_kinds:
+        alternatives.append("- break: break from the surrounding Python loop.\n")
+
+    if "continue" in allowed_kinds:
+        alternatives.append("- continue: continue to the next loop iteration.\n")
+
+    if "raise" in allowed_kinds:
+        raise_text = "- raise: raise a Python exception.\n  raise_message: required.\n"
+        if raise_error_type_binding_names:
+            error_type_names_text = ", ".join(f"`{name}`" for name in raise_error_type_binding_names)
+            raise_text += f"  raise_error_type: optional, must be one of: {error_type_names_text}.\n"
+        alternatives.append(raise_text)
+
+    if alternatives:
+        sections.append("\nAlternatives (use only when the program text explicitly requires it):\n")
+        sections.extend(alternatives)
+
+    return "".join(sections)