nighthawk-python 0.1.0__py3-none-any.whl

nighthawk/runtime/step_executor.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+from typing import Any, Protocol, cast, runtime_checkable
+
+from pydantic import TypeAdapter
+from pydantic_ai import Agent, StructuredDict
+from pydantic_ai.toolsets.function import FunctionToolset
+
+from ..configuration import StepExecutorConfiguration
+from ..errors import ExecutionError
+from ..tools.execution import ToolResultWrapperToolset
+from ..tools.registry import get_visible_tools
+from .async_bridge import run_coroutine_synchronously
+from .prompt import build_user_prompt, extract_references_and_program
+from .scoping import (
+    RUN_ID,
+    SCOPE_ID,
+    STEP_ID,
+    get_execution_context,
+    get_system_prompt_suffix_fragments,
+    scope,
+    span,
+)
+from .step_context import (
+    _MISSING,
+    StepContext,
+    ToolResultRenderingPolicy,
+    resolve_name_in_step_context,
+    step_context_scope,
+)
+from .step_contract import (
+    STEP_KINDS,
+    StepFinalResult,
+    StepKind,
+    StepOutcome,
+    build_step_json_schema,
+    build_step_system_prompt_suffix_fragment,
+)
+
+
+@runtime_checkable
+class AsyncExecutionAgent(Protocol):
+    """Protocol for agents that provide async step execution via ``run``."""
+
+    async def run(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+@runtime_checkable
+class SyncExecutionAgent(Protocol):
+    """Protocol for agents that provide sync step execution via ``run_sync``."""
+
+    def run_sync(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+type StepExecutionAgent = AsyncExecutionAgent | SyncExecutionAgent
+
+
+@runtime_checkable
+class SyncStepExecutor(Protocol):
+    """Step executor that provides synchronous execution."""
+
+    def run_step(
+        self,
+        *,
+        processed_natural_program: str,
+        step_context: StepContext,
+        binding_names: list[str],
+        allowed_step_kinds: tuple[str, ...],
+    ) -> tuple[StepOutcome, dict[str, object]]: ...
+
+
+@runtime_checkable
+class AsyncStepExecutor(Protocol):
+    """Step executor that provides asynchronous execution."""
+
+    async def run_step_async(
+        self,
+        *,
+        processed_natural_program: str,
+        step_context: StepContext,
+        binding_names: list[str],
+        allowed_step_kinds: tuple[str, ...],
+    ) -> tuple[StepOutcome, dict[str, object]]: ...
+
+
+type StepExecutor = SyncStepExecutor | AsyncStepExecutor
+
+
+def _new_agent_step_executor(
+    configuration: StepExecutorConfiguration,
+) -> StepExecutionAgent:
+    model_identifier = configuration.model
+    provider, provider_model_name = model_identifier.split(":", 1)
+
+    match provider:
+        case "claude-code-sdk":
+            from ..backends.claude_code_sdk import ClaudeCodeSdkModel
+
+            model: object = ClaudeCodeSdkModel(model_name=(provider_model_name if provider_model_name != "default" else None))
+        case "claude-code-cli":
+            from ..backends.claude_code_cli import ClaudeCodeCliModel
+
+            model = ClaudeCodeCliModel(model_name=(provider_model_name if provider_model_name != "default" else None))
+        case "codex":
+            from ..backends.codex import CodexModel
+
+            model = CodexModel(model_name=(provider_model_name if provider_model_name != "default" else None))
+        case _:
+            model = model_identifier
+
+    constructor_arguments: dict[str, Any] = {}
+    if configuration.model_settings is not None:
+        constructor_arguments["model_settings"] = configuration.model_settings
+
+    agent = Agent(
+        model=model,
+        output_type=StepFinalResult,
+        deps_type=StepContext,
+        system_prompt=configuration.prompts.step_system_prompt_template,
+        **constructor_arguments,
+    )
+
+    @agent.system_prompt(dynamic=True)
+    def _system_prompt_suffixes() -> str | None:  # pyright: ignore[reportUnusedFunction]
+        suffix_fragments = (
+            *configuration.system_prompt_suffix_fragments,
+            *get_system_prompt_suffix_fragments(),
+        )
+        if not suffix_fragments:
+            return None
+        return "\n\n".join(suffix_fragments)
+
+    return agent
+
+
+class AgentStepExecutor:
+    """Step executor that delegates Natural block execution to a Pydantic AI agent.
+
+    Attributes:
+        configuration: The step executor configuration.
+        agent: The underlying agent instance. If not provided, one is created
+            from the configuration.
+        token_encoding: The tiktoken encoding resolved from the configuration.
+        tool_result_rendering_policy: Policy for rendering tool results.
+        agent_is_managed: Whether the agent was created internally from
+            the configuration (True) or provided externally (False).
+    """
+
+    def __init__(
+        self,
+        configuration: StepExecutorConfiguration | None = None,
+        agent: StepExecutionAgent | None = None,
+    ) -> None:
+        self.configuration = configuration or StepExecutorConfiguration()
+        self.agent_is_managed = agent is None
+        self.agent = agent if agent is not None else _new_agent_step_executor(self.configuration)
+        self.token_encoding = self.configuration.resolve_token_encoding()
+        self.tool_result_rendering_policy = ToolResultRenderingPolicy(
+            tokenizer_encoding_name=self.token_encoding.name,
+            tool_result_max_tokens=(self.configuration.context_limits.tool_result_max_tokens),
+            json_renderer_style=self.configuration.json_renderer_style,
+        )
+
+    @classmethod
+    def from_agent(
+        cls,
+        *,
+        agent: StepExecutionAgent,
+        configuration: StepExecutorConfiguration | None = None,
+    ) -> AgentStepExecutor:
+        """Create an executor wrapping an existing agent.
+
+        Args:
+            agent: A pre-configured agent to use for step execution.
+            configuration: Optional configuration. Defaults to
+                StepExecutorConfiguration().
+        """
+        return cls(configuration=configuration, agent=agent)
+
+    @classmethod
+    def from_configuration(
+        cls,
+        *,
+        configuration: StepExecutorConfiguration,
+    ) -> AgentStepExecutor:
+        """Create an executor from a configuration, building a managed agent internally."""
+        return cls(configuration=configuration)
+
+    async def _run_agent(
+        self,
+        *,
+        user_prompt: str,
+        step_context: StepContext,
+        toolset: ToolResultWrapperToolset,
+        structured_output_type: object,
+    ) -> Any:
+        if self.agent is None:
+            raise ExecutionError("AgentStepExecutor.agent is not initialized")
+
+        if isinstance(self.agent, AsyncExecutionAgent):
+            return await self.agent.run(
+                user_prompt,
+                deps=step_context,
+                toolsets=[toolset],
+                output_type=structured_output_type,
+            )
+
+        if isinstance(self.agent, SyncExecutionAgent):
+            return self.agent.run_sync(
+                user_prompt,
+                deps=step_context,
+                toolsets=[toolset],
+                output_type=structured_output_type,
+            )
+
+        raise ExecutionError("AgentStepExecutor requires an agent with run(...) or run_sync(...)")
+
+    def _build_structured_output_and_prompt_fragment(
+        self,
+        *,
+        processed_natural_program: str,
+        step_context: StepContext,
+        allowed_step_kinds: tuple[str, ...],
+    ) -> tuple[object, str]:
+        """Build the structured output type and system prompt fragment for a step."""
+        unknown_kinds = set(allowed_step_kinds).difference(STEP_KINDS)
+        if unknown_kinds:
+            raise ExecutionError(f"Internal error: allowed_step_kinds contains unknown kinds: {tuple(sorted(unknown_kinds))}")
+
+        allowed_kinds_deduplicated = tuple(dict.fromkeys(allowed_step_kinds))
+        allowed_kinds_typed = cast(tuple[StepKind, ...], allowed_kinds_deduplicated)
+
+        referenced_names, _ = extract_references_and_program(processed_natural_program)
+
+        error_type_candidate_names: set[str] = set(referenced_names)
+        for name, value in step_context.step_locals.items():
+            if isinstance(value, type) and issubclass(value, BaseException) and value.__name__ == name:
+                error_type_candidate_names.add(name)
+
+        error_type_binding_name_list: list[str] = []
+        for name in sorted(error_type_candidate_names):
+            value = resolve_name_in_step_context(step_context, name)
+            if value is _MISSING:
+                continue
+            if not isinstance(value, type) or not issubclass(value, BaseException) or value.__name__ != name:
+                continue
+            error_type_binding_name_list.append(name)
+
+        raise_error_type_binding_names = tuple(error_type_binding_name_list)
+
+        step_system_prompt_fragment = build_step_system_prompt_suffix_fragment(
+            allowed_kinds=allowed_kinds_typed,
+            raise_error_type_binding_names=raise_error_type_binding_names,
+        )
+
+        outcome_json_schema = build_step_json_schema(
+            allowed_kinds=allowed_kinds_typed,
+            raise_error_type_binding_names=raise_error_type_binding_names,
+        )
+        structured_output_type = StructuredDict(outcome_json_schema, name="StepFinalResult")
+
+        return structured_output_type, step_system_prompt_fragment
+
+    def _parse_agent_result(
+        self,
+        result: Any,
+    ) -> StepOutcome:
+        """Parse the agent result into a StepOutcome."""
+        try:
+            raw_output = result.output
+            if isinstance(raw_output, StepFinalResult):
+                return raw_output.result
+            if isinstance(raw_output, dict) and "result" in raw_output:
+                return TypeAdapter(StepOutcome).validate_python(raw_output["result"])
+            return TypeAdapter(StepOutcome).validate_python(raw_output)
+        except Exception as e:
+            raise ExecutionError(f"Step produced invalid step outcome: {e}") from e
+
+    def _extract_bindings(
+        self,
+        *,
+        binding_names: list[str],
+        step_context: StepContext,
+    ) -> dict[str, object]:
+        """Extract committed bindings from the step context."""
+        bindings: dict[str, object] = {}
+        for name in binding_names:
+            if name in step_context.assigned_binding_names:
+                bindings[name] = step_context.step_locals[name]
+        return bindings
+
+    async def run_step_async(
+        self,
+        *,
+        processed_natural_program: str,
+        step_context: StepContext,
+        binding_names: list[str],
+        allowed_step_kinds: tuple[str, ...],
+    ) -> tuple[StepOutcome, dict[str, object]]:
+        if step_context.tool_result_rendering_policy is None:
+            step_context.tool_result_rendering_policy = self.tool_result_rendering_policy
+
+        user_prompt = build_user_prompt(
+            processed_natural_program=processed_natural_program,
+            step_context=step_context,
+            configuration=self.configuration,
+        )
+
+        visible_tool_list = get_visible_tools()
+        toolset = ToolResultWrapperToolset(FunctionToolset(visible_tool_list))
+
+        structured_output_type, step_system_prompt_fragment = self._build_structured_output_and_prompt_fragment(
+            processed_natural_program=processed_natural_program,
+            step_context=step_context,
+            allowed_step_kinds=allowed_step_kinds,
+        )
+
+        with scope(system_prompt_suffix_fragment=step_system_prompt_fragment):
+            execution_context = get_execution_context()
+            with (
+                span(
+                    "nighthawk.step_executor",
+                    **{
+                        RUN_ID: execution_context.run_id,
+                        SCOPE_ID: execution_context.scope_id,
+                        STEP_ID: step_context.step_id,
+                    },
+                ),
+                step_context_scope(step_context),
+            ):
+                result = await self._run_agent(
+                    user_prompt=user_prompt,
+                    step_context=step_context,
+                    toolset=toolset,
+                    structured_output_type=structured_output_type,
+                )
+
+        step_outcome = self._parse_agent_result(result)
+        bindings = self._extract_bindings(binding_names=binding_names, step_context=step_context)
+        return step_outcome, bindings
+
+    def run_step(
+        self,
+        *,
+        processed_natural_program: str,
+        step_context: StepContext,
+        binding_names: list[str],
+        allowed_step_kinds: tuple[str, ...],
+    ) -> tuple[StepOutcome, dict[str, object]]:
+        return cast(
+            tuple[StepOutcome, dict[str, object]],
+            run_coroutine_synchronously(
+                lambda: self.run_step_async(
+                    processed_natural_program=processed_natural_program,
+                    step_context=step_context,
+                    binding_names=binding_names,
+                    allowed_step_kinds=allowed_step_kinds,
+                )
+            ),
+        )

nighthawk/runtime/tool_calls.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import json
+import uuid
+from collections.abc import Awaitable, Callable, Iterator
+from contextlib import contextmanager
+from typing import Any
+
+from pydantic_ai import RunContext
+from pydantic_ai._instrumentation import InstrumentationNames
+
+
+def generate_tool_call_id() -> str:
+    return str(uuid.uuid4())
+
+
+def _resolve_instrumentation_names(*, run_context: RunContext[Any]) -> InstrumentationNames:
+    return InstrumentationNames.for_version(run_context.instrumentation_version)
+
+
+def _resolve_trace_include_content(*, run_context: RunContext[Any]) -> bool:
+    return run_context.trace_include_content
+
+
+def _build_tool_span_attributes(
+    *,
+    tool_name: str,
+    tool_call_id: str | None,
+    arguments: dict[str, Any],
+    instrumentation_names: InstrumentationNames,
+    include_content: bool,
+) -> dict[str, Any]:
+    attributes: dict[str, Any] = {
+        "gen_ai.tool.name": tool_name,
+        "logfire.msg": f"running tool: {tool_name}",
+    }
+
+    if tool_call_id is not None:
+        attributes["gen_ai.tool.call.id"] = tool_call_id
+
+    if include_content:
+        attributes[instrumentation_names.tool_arguments_attr] = json.dumps(arguments, default=str)
+        attributes["logfire.json_schema"] = json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    instrumentation_names.tool_arguments_attr: {"type": "object"},
+                    instrumentation_names.tool_result_attr: {"type": "object"},
+                    "gen_ai.tool.name": {},
+                    "gen_ai.tool.call.id": {},
+                },
+            }
+        )
+
+    return attributes
+
+
+@contextmanager
+def _start_tool_span(
+    *,
+    tool_name: str,
+    attributes: dict[str, Any],
+    instrumentation_names: InstrumentationNames,
+    run_context: RunContext[Any],
+) -> Iterator[Any]:
+    span_name = instrumentation_names.get_tool_span_name(tool_name)
+    with run_context.tracer.start_as_current_span(span_name, attributes=attributes) as span:
+        yield span
+
+
+async def run_tool_instrumented(
+    *,
+    tool_name: str,
+    arguments: dict[str, Any],
+    call: Callable[[], Awaitable[str]],
+    run_context: RunContext[Any],
+    tool_call_id: str | None,
+) -> str:
+    instrumentation_names = _resolve_instrumentation_names(run_context=run_context)
+    include_content = _resolve_trace_include_content(run_context=run_context)
+
+    span_attributes = _build_tool_span_attributes(
+        tool_name=tool_name,
+        tool_call_id=tool_call_id,
+        arguments=arguments,
+        instrumentation_names=instrumentation_names,
+        include_content=include_content,
+    )
+
+    with _start_tool_span(
+        tool_name=tool_name,
+        attributes=span_attributes,
+        instrumentation_names=instrumentation_names,
+        run_context=run_context,
+    ) as span:
+        result_text = await call()
+        if include_content and span.is_recording():
+            span.set_attribute(instrumentation_names.tool_result_attr, result_text)
+        return result_text

nighthawk/tools/assignment.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import ast
+import inspect
+import types
+from typing import Any, NoReturn
+
+from pydantic import BaseModel, TypeAdapter
+
+from ..errors import ToolEvaluationError
+from ..identifier_path import parse_identifier_path
+from ..json_renderer import to_jsonable_value
+from ..runtime.async_bridge import run_awaitable_value_synchronously
+from ..runtime.step_context import StepContext
+from .contracts import ToolBoundaryError
+
+
+def _compile_expression(expression: str) -> types.CodeType:
+    """Compile a Python expression with top-level await support."""
+    return compile(
+        expression,
+        "<nighthawk-eval>",
+        "eval",
+        flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+    )
+
+
+async def eval_expression_async(step_context: StepContext, expression: str) -> object:
+    """Evaluate a Python expression inside the step execution environment (async canonical).
+
+    Security note: This uses ``eval()`` with ``ast.PyCF_ALLOW_TOP_LEVEL_AWAIT``
+    to execute LLM-generated expressions against ``step_globals`` and ``step_locals``.
+    Natural DSL sources are trusted, repository-managed assets.
+    Do not wire untrusted user input into expressions evaluated here.
+    """
+    try:
+        compiled_expression = _compile_expression(expression)
+        value = eval(compiled_expression, step_context.step_globals, step_context.step_locals)
+        if inspect.isawaitable(value):
+            return await value
+        return value
+    except Exception as e:
+        raise ToolEvaluationError(str(e)) from e
+
+
+def eval_expression(step_context: StepContext, expression: str) -> object:
+    """Evaluate a Python expression inside the step execution environment (sync wrapper)."""
+    try:
+        compiled_expression = _compile_expression(expression)
+        value = eval(compiled_expression, step_context.step_globals, step_context.step_locals)
+        return run_awaitable_value_synchronously(value)
+    except Exception as e:
+        raise ToolEvaluationError(str(e)) from e
+
+
+def _raise_invalid_input(*, message: str, guidance: str) -> NoReturn:
+    raise ToolBoundaryError(kind="invalid_input", message=message, guidance=guidance)
+
+
+def _raise_resolution(*, message: str, guidance: str) -> NoReturn:
+    raise ToolBoundaryError(kind="resolution", message=message, guidance=guidance)
+
+
+def _raise_execution(*, message: str, guidance: str) -> NoReturn:
+    raise ToolBoundaryError(kind="execution", message=message, guidance=guidance)
+
+
+def _get_pydantic_field_type(model: BaseModel, field_name: str) -> object | None:
+    model_fields = getattr(type(model), "model_fields", None)
+    if model_fields is None:
+        return None
+    field = model_fields.get(field_name)
+    if field is None:
+        return None
+    return field.annotation
+
+
+def _assign_value_to_target_path(
+    *,
+    step_context: StepContext,
+    target_path: str,
+    parsed_target_path: tuple[str, ...],
+    value: object,
+) -> dict[str, Any]:
+    if len(parsed_target_path) == 1:
+        name = parsed_target_path[0]
+
+        if name in step_context.read_binding_names:
+            _raise_invalid_input(
+                message=f"Cannot rebind read binding '{name}' with nh_assign.",
+                guidance=f"'{name}' is a read binding (<{name}>). To mutate it in-place, use nh_exec (e.g. {name}.update(...)). To rebind, declare it as a write binding (<:{name}>).",
+            )
+
+        expected_type = step_context.binding_name_to_type.get(name)
+
+        if expected_type is not None:
+            try:
+                adapted = TypeAdapter(expected_type)
+                value = adapted.validate_python(value)
+            except Exception as e:
+                _raise_invalid_input(
+                    message=str(e),
+                    guidance="Fix the value to match the expected type and retry.",
+                )
+
+        step_context.record_assignment(name, value)
+
+        return {
+            "target_path": target_path,
+            "step_locals_revision": step_context.step_locals_revision,
+            "updates": [{"path": target_path, "value": to_jsonable_value(value)}],
+        }
+
+    root_name = parsed_target_path[0]
+    attribute_path = parsed_target_path[1:]
+
+    if root_name not in step_context.step_locals:
+        _raise_resolution(
+            message=f"Unknown root name: {root_name}",
+            guidance="Fix the target path so the referenced root name exists, then retry.",
+        )
+    root_object = step_context.step_locals[root_name]
+
+    current_object = root_object
+    for attribute in attribute_path[:-1]:
+        try:
+            current_object = getattr(current_object, attribute)
+        except Exception as e:
+            _raise_resolution(
+                message=f"Failed to resolve attribute {attribute!r}: {e}",
+                guidance="Fix the target path so the referenced attributes exist, then retry.",
+            )
+
+    final_attribute = attribute_path[-1]
+
+    expected_type: object | None = None
+    if isinstance(current_object, BaseModel):
+        expected_type = _get_pydantic_field_type(current_object, final_attribute)
+        if expected_type is None:
+            _raise_invalid_input(
+                message=f"Unknown field on {type(current_object).__name__}: {final_attribute}",
+                guidance="Fix the target path so the referenced field exists, then retry.",
+            )
+
+    if expected_type is not None:
+        try:
+            adapted = TypeAdapter(expected_type)
+            value = adapted.validate_python(value)
+        except Exception as e:
+            _raise_invalid_input(
+                message=str(e),
+                guidance="Fix the value to match the expected type and retry.",
+            )
+
+    try:
+        setattr(current_object, final_attribute, value)
+    except Exception as e:
+        _raise_resolution(
+            message=f"Failed to set attribute {final_attribute!r}: {e}",
+            guidance="Fix the target path so the referenced attributes are assignable, then retry.",
+        )
+
+    # Dotted mutation bypasses record_assignment because commit selection is
+    # controlled only by <:name> bindings (top-level names).  See design.md
+    # Section 8.3 "Commit and mutation notes" for the distinction.
+    step_context.step_locals_revision += 1
+
+    return {
+        "target_path": target_path,
+        "step_locals_revision": step_context.step_locals_revision,
+        "updates": [{"path": target_path, "value": to_jsonable_value(value)}],
+    }
+
+
+def _resolve_value_for_assignment(step_context: StepContext, expression: str) -> object:
+    try:
+        return eval_expression(step_context, expression)
+    except Exception as e:
+        _raise_execution(
+            message=str(e),
+            guidance="Fix the expression and retry.",
+        )
+
+
+async def _resolve_value_for_assignment_async(step_context: StepContext, expression: str) -> object:
+    try:
+        return await eval_expression_async(step_context, expression)
+    except Exception as e:
+        _raise_execution(
+            message=str(e),
+            guidance="Fix the expression and retry.",
+        )
+
+
+def _validated_target_path(target_path: str) -> tuple[str, ...]:
+    parsed = parse_identifier_path(target_path)
+    if parsed is None:
+        _raise_invalid_input(
+            message="Invalid target_path; expected name(.field)* with ASCII identifiers",
+            guidance="Fix target_path to match name(.field)* with ASCII identifiers and retry.",
+        )
+    return parsed
+
+
+def assign_tool(
+    step_context: StepContext,
+    target_path: str,
+    expression: str,
+) -> dict[str, Any]:
+    """Assign a computed value to a dotted target_path.
+
+    Target grammar:
+    - target_path := name ("." field)*
+
+    Notes:
+    - Any segment starting with "__" is forbidden.
+    - On success, returns a JSON-serializable payload with keys `target_path`, `step_locals_revision`, and `updates`.
+    - On failure, raises ToolBoundaryError.
+    - The operation is atomic: on any failure, no updates are performed.
+    """
+    parsed_target_path = _validated_target_path(target_path)
+    value = _resolve_value_for_assignment(step_context, expression)
+
+    return _assign_value_to_target_path(
+        step_context=step_context,
+        target_path=target_path,
+        parsed_target_path=parsed_target_path,
+        value=value,
+    )
+
+
+async def assign_tool_async(
+    step_context: StepContext,
+    target_path: str,
+    expression: str,
+) -> dict[str, Any]:
+    """Async version of assign_tool."""
+    parsed_target_path = _validated_target_path(target_path)
+    value = await _resolve_value_for_assignment_async(step_context, expression)
+
+    return _assign_value_to_target_path(
+        step_context=step_context,
+        target_path=target_path,
+        parsed_target_path=parsed_target_path,
+        value=value,
+    )