From b629b4d56bf848839c8b7ca1bdc5f25af2387a0d Mon Sep 17 00:00:00 2001 From: Grzegorz Ziolkowski Date: Thu, 23 Jul 2026 13:38:20 +0200 Subject: [PATCH 1/3] AI: Add automatic ability resolution to the WP AI Client. Introduces using_ability_resolution() on WP_AI_Client_Prompt_Builder. When enabled, the text generation methods run a resolution loop: each round executes the ability function calls requested by the model, appends the results to the conversation, and requests a follow-up response, until the model produces a final answer, requests an unknown function, or the maximum number of rounds is reached. Also adds the wp_ai_client_ability_resolution_defaults and wp_ai_client_pre_resolve_ability_call filters and the wp_ai_client_ability_call_resolved action. See #64865. Co-Authored-By: Claude Fable 5 --- ...wp-ai-client-ability-function-resolver.php | 72 +- .../class-wp-ai-client-prompt-builder.php | 473 +++++++++++- ...wp-ai-client-mock-model-creation-trait.php | 78 ++ .../wpAiClientAbilityFunctionResolver.php | 117 +++ .../ai-client/wpAiClientAbilityResolution.php | 731 ++++++++++++++++++ 5 files changed, 1442 insertions(+), 29 deletions(-) create mode 100644 tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php diff --git a/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php b/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php index 596656d36f516..0a755539d4340 100644 --- a/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php +++ b/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php @@ -75,7 +75,19 @@ public function is_ability_call( FunctionCall $call ): bool { return false; } - return str_starts_with( $name, self::ABILITY_PREFIX ); + return self::is_ability_function_name( $name ); + } + + /** + * Checks if a function name refers to an ability. + * + * @since 7.2.0 + * + * @param string $function_name The function name to check. + * @return bool True if the function name refers to an ability, false otherwise. + */ + public static function is_ability_function_name( string $function_name ): bool { + return str_starts_with( $function_name, self::ABILITY_PREFIX ); } /** @@ -91,6 +103,34 @@ public function is_ability_call( FunctionCall $call ): bool { * @return FunctionResponse The response from executing the ability. */ public function execute_ability( FunctionCall $call ): FunctionResponse { + $response = $this->resolve_ability_call( $call ); + + /** + * Fires after an ability function call has been resolved. + * + * The response may represent a successful execution or an error, for + * example when the ability is not allowed or its execution failed. + * Useful for logging and auditing ability usage by AI models. + * + * @since 7.2.0 + * + * @param FunctionCall $call The function call requested by the model. + * @param FunctionResponse $response The response that will be sent back to the model. + */ + do_action( 'wp_ai_client_ability_call_resolved', $call, $response ); + + return $response; + } + + /** + * Resolves a function call into a function response. + * + * @since 7.2.0 + * + * @param FunctionCall $call The function call to resolve. + * @return FunctionResponse The response from resolving the call. + */ + private function resolve_ability_call( FunctionCall $call ): FunctionResponse { $function_name = $call->getName() ?? 'unknown'; $function_id = $call->getId() ?? 'unknown'; @@ -133,6 +173,36 @@ public function execute_ability( FunctionCall $call ): FunctionResponse { ); } + /** + * Filters whether to short-circuit the execution of an ability function call. + * + * Returning a FunctionResponse skips the ability execution and sends that + * response back to the model instead. Returning a WP_Error skips the + * execution and sends an error response back to the model. This allows + * vetoing individual calls, for example based on ability annotations. + * + * @since 7.2.0 + * + * @param FunctionResponse|WP_Error|null $pre A response to short-circuit with, or null to execute the ability. Default null. + * @param FunctionCall $call The function call requested by the model. + * @param WP_Ability $ability The ability that is about to be executed. + */ + $pre = apply_filters( 'wp_ai_client_pre_resolve_ability_call', null, $call, $ability ); + if ( $pre instanceof FunctionResponse ) { + return $pre; + } + if ( is_wp_error( $pre ) ) { + return new FunctionResponse( + $function_id, + $function_name, + array( + 'error' => $pre->get_error_message(), + 'code' => $pre->get_error_code(), + 'data' => $pre->get_error_data(), + ) + ); + } + $args = $call->getArgs(); $result = $ability->execute( ! empty( $args ) ? $args : null ); diff --git a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php index a64957fe73157..84024e3c35e82 100644 --- a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php +++ b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php @@ -11,6 +11,8 @@ use WordPress\AiClient\Builders\PromptBuilder; use WordPress\AiClient\Common\Exception\InvalidArgumentException; use WordPress\AiClient\Common\Exception\TokenLimitReachedException; +use WordPress\AiClient\Events\AfterGenerateResultEvent; +use WordPress\AiClient\Events\BeforeGenerateResultEvent; use WordPress\AiClient\Files\DTO\File; use WordPress\AiClient\Files\Enums\FileTypeEnum; use WordPress\AiClient\Files\Enums\MediaOrientationEnum; @@ -24,8 +26,11 @@ use WordPress\AiClient\Providers\Models\Contracts\ModelInterface; use WordPress\AiClient\Providers\Models\DTO\ModelConfig; use WordPress\AiClient\Providers\Models\Enums\CapabilityEnum; +use WordPress\AiClient\Providers\Models\TextGeneration\Contracts\TextGenerationModelInterface; use WordPress\AiClient\Providers\ProviderRegistry; use WordPress\AiClient\Results\DTO\GenerativeAiResult; +use WordPress\AiClient\Results\DTO\TokenUsage; +use WordPress\AiClient\Tools\DTO\FunctionCall; use WordPress\AiClient\Tools\DTO\FunctionDeclaration; use WordPress\AiClient\Tools\DTO\FunctionResponse; use WordPress\AiClient\Tools\DTO\WebSearch; @@ -122,6 +127,14 @@ class WP_AI_Client_Prompt_Builder { */ private ?WP_Error $error = null; + /** + * Options for automatic ability resolution, or null when disabled. + * + * @since 7.2.0 + * @var array{max_iterations: int}|null + */ + private ?array $ability_resolution_options = null; + /** * List of methods that generate a result from the prompt. * @@ -277,12 +290,86 @@ public function using_abilities( ...$abilities ): self { return $this; } + /** + * Enables automatic resolution of ability function calls. + * + * When enabled, the text generation methods run a resolution loop instead of + * a single request. Each round executes the ability calls requested by the + * model, appends the results to the conversation, and requests a follow-up + * response. The loop ends when the model produces a response without ability + * calls, when it requests a function that is not a registered ability, or + * when the maximum number of rounds is reached. + * + * Only abilities that were exposed to the model as function declarations, + * typically with {@see self::using_abilities()}, can be executed. + * Resolution follows the first response candidate and supports the + * generate_text_result() and generate_text() methods. Details about the loop + * are exposed under the `ability_resolution` key of the additional data of + * the final result. + * + * @since 7.2.0 + * + * @param array $options { + * Optional. Options controlling the resolution loop. + * + * @type int $max_iterations Maximum number of resolution rounds. Each round executes + * the ability calls from one model response and requests a + * follow-up response. Default 5. + * } + * @return self The current instance for method chaining. + */ + public function using_ability_resolution( array $options = array() ): self { + $defaults = array( + 'max_iterations' => 5, + ); + + /** + * Filters the default options for automatic ability resolution. + * + * @since 7.2.0 + * + * @param array $defaults Default options. See {@see WP_AI_Client_Prompt_Builder::using_ability_resolution()} + * for the supported keys. + */ + $filtered_defaults = apply_filters( 'wp_ai_client_ability_resolution_defaults', $defaults ); + if ( is_array( $filtered_defaults ) ) { + $defaults = wp_parse_args( $filtered_defaults, $defaults ); + } + + // Guard against invalid filtered defaults. + if ( ! is_int( $defaults['max_iterations'] ) || $defaults['max_iterations'] < 1 ) { + $defaults['max_iterations'] = 5; + } + + $options = wp_parse_args( $options, $defaults ); + + if ( ! is_int( $options['max_iterations'] ) || $options['max_iterations'] < 1 ) { + _doing_it_wrong( + __METHOD__, + sprintf( + /* translators: %s: max_iterations */ + __( 'The %s option must be a positive integer.' ), + 'max_iterations' + ), + '7.2.0' + ); + $options['max_iterations'] = $defaults['max_iterations']; + } + + $this->ability_resolution_options = array( + 'max_iterations' => $options['max_iterations'], + ); + + return $this; + } + /** * Magic method to proxy snake_case method calls to their PHP AI Client camelCase counterparts. * * This allows WordPress developers to use snake_case naming conventions. It catches * any exceptions thrown, stores them, and returns a WP_Error when a terminate method - * is called. + * is called. When automatic ability resolution is enabled, the supported text + * generation methods run the resolution loop instead of a single request. * * @since 7.0.0 * @@ -291,6 +378,37 @@ public function using_abilities( ...$abilities ): self { * @return mixed The result of the method call. */ public function __call( string $name, array $arguments ) { + if ( null !== $this->ability_resolution_options && self::is_generating_method( $name ) ) { + if ( 'generate_text_result' === $name || 'generate_text' === $name ) { + return $this->generate_with_ability_resolution( $name ); + } + + _doing_it_wrong( + __METHOD__, + sprintf( + /* translators: 1: generate_text_result, 2: generate_text, 3: the method that was called. */ + __( 'Automatic ability resolution supports only the %1$s and %2$s methods. The %3$s method runs without it.' ), + 'generate_text_result()', + 'generate_text()', + '' . esc_html( $name ) . '()' + ), + '7.2.0' + ); + } + + return $this->call_builder( $name, $arguments ); + } + + /** + * Proxies a method call to the wrapped prompt builder with WordPress-specific guards. + * + * @since 7.2.0 + * + * @param string $name The method name in snake_case. + * @param array $arguments The method arguments. + * @return mixed The result of the method call. + */ + private function call_builder( string $name, array $arguments ) { /* * If an error occurred in a previous method call, either return the error for terminate methods, * or return the same instance for other methods to maintain the fluent interface. @@ -307,39 +425,16 @@ public function __call( string $name, array $arguments ) { // Check if the prompt should be prevented for is_supported* and generate_*/convert_text_to_speech* methods. if ( self::is_support_check_method( $name ) || self::is_generating_method( $name ) ) { - // If AI is not supported, then there's no need to apply the filter as the prompt will be prevented anyway. - $is_ai_disabled = ! wp_supports_ai(); - $prevent = $is_ai_disabled; - if ( ! $prevent ) { - /** - * Filters whether to prevent the prompt from being executed. - * - * @since 7.0.0 - * - * @param bool $prevent Whether to prevent the prompt. Default false. - * @param WP_AI_Client_Prompt_Builder $builder A clone of the prompt builder instance (read-only). - */ - $prevent = (bool) apply_filters( 'wp_ai_client_prevent_prompt', false, clone $this ); - } + $prevented = $this->get_prompt_prevented_error(); - if ( $prevent ) { + if ( null !== $prevented ) { // For is_supported* methods, return false. if ( self::is_support_check_method( $name ) ) { return false; } - $error_message = $is_ai_disabled - ? __( 'AI features are not supported in this environment.' ) - : __( 'Prompt execution was prevented by a filter.' ); - - // For generate_* and convert_text_to_speech* methods, create a WP_Error. - $this->error = new WP_Error( - 'prompt_prevented', - $error_message, - array( - 'status' => 503, - ) - ); + // For generate_* and convert_text_to_speech* methods, store the WP_Error. + $this->error = $prevented; if ( self::is_generating_method( $name ) ) { return $this->error; @@ -368,6 +463,328 @@ public function __call( string $name, array $arguments ) { } } + /** + * Checks whether the prompt is prevented from being executed. + * + * @since 7.2.0 + * + * @return WP_Error|null A WP_Error when the prompt is prevented, null otherwise. + */ + private function get_prompt_prevented_error(): ?WP_Error { + // If AI is not supported, then there's no need to apply the filter as the prompt will be prevented anyway. + $is_ai_disabled = ! wp_supports_ai(); + $prevent = $is_ai_disabled; + if ( ! $prevent ) { + /** + * Filters whether to prevent the prompt from being executed. + * + * @since 7.0.0 + * + * @param bool $prevent Whether to prevent the prompt. Default false. + * @param WP_AI_Client_Prompt_Builder $builder A clone of the prompt builder instance (read-only). + */ + $prevent = (bool) apply_filters( 'wp_ai_client_prevent_prompt', false, clone $this ); + } + + if ( ! $prevent ) { + return null; + } + + $error_message = $is_ai_disabled + ? __( 'AI features are not supported in this environment.' ) + : __( 'Prompt execution was prevented by a filter.' ); + + return new WP_Error( + 'prompt_prevented', + $error_message, + array( + 'status' => 503, + ) + ); + } + + /** + * Generates a text result while automatically resolving ability function calls. + * + * Runs the resolution loop: each round executes the ability calls requested + * by the model, appends the results to the conversation, and requests a + * follow-up response. See {@see self::using_ability_resolution()} for the + * termination conditions. + * + * @since 7.2.0 + * + * @param string $method Either 'generate_text_result' or 'generate_text'. + * @return GenerativeAiResult|string|WP_Error The final result, the final text, or a WP_Error on failure. + */ + private function generate_with_ability_resolution( string $method ) { + $options = $this->ability_resolution_options; + + /* + * The PHP AI Client prompt builder does not expose its message list, nor + * a way to append messages to it. The first request therefore captures + * the sent messages and the resolved model from the lifecycle event that + * the builder dispatches. Later rounds call the captured model directly + * with an extended copy of that transcript. A message append method in + * the PHP AI Client would simplify this. + */ + $captured = null; + $capture = static function ( $event ) use ( &$captured ) { + if ( null === $captured && $event instanceof BeforeGenerateResultEvent ) { + $captured = $event; + } + }; + + add_action( 'wp_ai_client_before_generate_result', $capture ); + $result = $this->call_builder( 'generate_text_result', array() ); + remove_action( 'wp_ai_client_before_generate_result', $capture ); + + if ( is_wp_error( $result ) ) { + return $result; + } + + if ( null === $captured || ! $captured->getModel() instanceof TextGenerationModelInterface ) { + // Without the captured context the conversation cannot be continued. + return $this->to_generation_return_value( $result, $method ); + } + + $model = $captured->getModel(); + $capability = $captured->getCapability(); + $transcript = $captured->getMessages(); + $dispatcher = AiClient::getEventDispatcher(); + + /* + * The allow-list for execution is derived from the function declarations + * that were sent to the model. A response may name any function, so the + * resolver enforces that only explicitly exposed abilities can run. + */ + $ability_names = array(); + $declarations = $model->getConfig()->getFunctionDeclarations() ?? array(); + foreach ( $declarations as $declaration ) { + $function_name = $declaration->getName(); + if ( WP_AI_Client_Ability_Function_Resolver::is_ability_function_name( $function_name ) ) { + $ability_names[] = WP_AI_Client_Ability_Function_Resolver::function_name_to_ability_name( $function_name ); + } + } + + if ( empty( $ability_names ) ) { + _doing_it_wrong( + __METHOD__, + sprintf( + /* translators: 1: using_ability_resolution, 2: using_abilities */ + __( '%1$s requires abilities registered with %2$s.' ), + 'using_ability_resolution()', + 'using_abilities()' + ), + '7.2.0' + ); + return $this->to_generation_return_value( $result, $method ); + } + + $resolver = new WP_AI_Client_Ability_Function_Resolver( ...$ability_names ); + + $rounds = 0; + $usage = $result->getTokenUsage(); + $resolved_calls = array(); + + while ( true ) { + $message = $result->toMessage(); + $calls = $this->get_function_calls( $message ); + + if ( empty( $calls ) ) { + $stop_reason = 'completed'; + break; + } + + $ability_calls = array_filter( $calls, array( $resolver, 'is_ability_call' ) ); + + if ( count( $ability_calls ) < count( $calls ) ) { + // The response requests functions that are not registered abilities. + // Hand the round back to the caller to resolve them. + $stop_reason = 'unresolved_function_calls'; + break; + } + + if ( $rounds >= $options['max_iterations'] ) { + $stop_reason = 'max_iterations'; + break; + } + + $responses = $resolver->execute_abilities( $message ); + + foreach ( $ability_calls as $call ) { + $resolved_calls[] = array( + 'id' => $call->getId(), + 'ability' => WP_AI_Client_Ability_Function_Resolver::function_name_to_ability_name( (string) $call->getName() ), + ); + } + + $transcript[] = $message; + $transcript[] = $responses; + ++$rounds; + + $prevented = $this->get_prompt_prevented_error(); + if ( null !== $prevented ) { + $this->error = $prevented; + return $this->error; + } + + if ( null !== $dispatcher ) { + $dispatcher->dispatch( new BeforeGenerateResultEvent( $transcript, $model, $capability ) ); + } + + try { + $result = $model->generateTextResult( $transcript ); + } catch ( Exception $e ) { + $this->error = $this->exception_to_wp_error( $e ); + return $this->error; + } + + if ( null !== $dispatcher ) { + $dispatcher->dispatch( new AfterGenerateResultEvent( $transcript, $model, $capability, $result ) ); + } + + $usage = $this->aggregate_token_usage( $usage, $result->getTokenUsage() ); + } + + return $this->finish_ability_resolution( $result, $method, $stop_reason, $rounds, $usage, $resolved_calls, $transcript ); + } + + /** + * Converts a result into the return value of the called generation method. + * + * Used when the resolution loop exits early with a plain result, so that + * generate_text() still returns a string or a WP_Error. + * + * @since 7.2.0 + * + * @param GenerativeAiResult $result The result to convert. + * @param string $method Either 'generate_text_result' or 'generate_text'. + * @return GenerativeAiResult|string|WP_Error The result, its text, or a WP_Error on failure. + */ + private function to_generation_return_value( GenerativeAiResult $result, string $method ) { + if ( 'generate_text' !== $method ) { + return $result; + } + + try { + return $result->toText(); + } catch ( Exception $e ) { + $this->error = $this->exception_to_wp_error( $e ); + return $this->error; + } + } + + /** + * Retrieves the function calls contained in a message. + * + * @since 7.2.0 + * + * @param Message $message The message to inspect. + * @return FunctionCall[] The function calls in the message. + */ + private function get_function_calls( Message $message ): array { + $calls = array(); + + foreach ( $message->getParts() as $part ) { + if ( $part->getType()->isFunctionCall() ) { + $call = $part->getFunctionCall(); + if ( $call instanceof FunctionCall ) { + $calls[] = $call; + } + } + } + + return $calls; + } + + /** + * Adds up two token usage objects. + * + * @since 7.2.0 + * + * @param TokenUsage $total The running total. + * @param TokenUsage $addition The usage to add. + * @return TokenUsage The combined token usage. + */ + private function aggregate_token_usage( TokenUsage $total, TokenUsage $addition ): TokenUsage { + $thought_tokens = null; + if ( null !== $total->getThoughtTokens() || null !== $addition->getThoughtTokens() ) { + $thought_tokens = (int) $total->getThoughtTokens() + (int) $addition->getThoughtTokens(); + } + + return new TokenUsage( + $total->getPromptTokens() + $addition->getPromptTokens(), + $total->getCompletionTokens() + $addition->getCompletionTokens(), + $total->getTotalTokens() + $addition->getTotalTokens(), + $thought_tokens + ); + } + + /** + * Builds the final value of an ability resolution loop. + * + * Rebuilds the result with the aggregated token usage and details about the + * loop under the `ability_resolution` key of the additional data. + * + * @since 7.2.0 + * + * @param GenerativeAiResult $result The result of the last round. + * @param string $method Either 'generate_text_result' or 'generate_text'. + * @param string $stop_reason Why the loop ended. One of 'completed', + * 'unresolved_function_calls', or 'max_iterations'. + * @param int $rounds Number of resolution rounds that ran. + * @param TokenUsage $usage Aggregated token usage across all rounds. + * @param array $resolved_calls The ability calls that were resolved. + * @param Message[] $transcript The conversation before the final response. + * @return GenerativeAiResult|string|WP_Error The final result or text, or a WP_Error on failure. + */ + private function finish_ability_resolution( GenerativeAiResult $result, string $method, string $stop_reason, int $rounds, TokenUsage $usage, array $resolved_calls, array $transcript ) { + $messages = $transcript; + $messages[] = $result->toMessage(); + + $additional_data = $result->getAdditionalData(); + $additional_data['ability_resolution'] = array( + 'rounds' => $rounds, + 'stop_reason' => $stop_reason, + 'resolved_calls' => $resolved_calls, + 'messages' => array_map( + static function ( Message $message ) { + return $message->toArray(); + }, + $messages + ), + ); + + $final = new GenerativeAiResult( + $result->getId(), + $result->getCandidates(), + $usage, + $result->getProviderMetadata(), + $result->getModelMetadata(), + $additional_data + ); + + if ( 'generate_text_result' === $method ) { + return $final; + } + + // generate_text() returns the plain text of the final answer. + if ( 'completed' !== $stop_reason ) { + $this->error = new WP_Error( + 'ability_resolution_incomplete', + __( 'The model did not produce a final answer within the ability resolution limits.' ), + array( + 'status' => 500, + 'stop_reason' => $stop_reason, + 'rounds' => $rounds, + ) + ); + return $this->error; + } + + return $this->to_generation_return_value( $final, $method ); + } + /** * Converts an exception into a WP_Error with a structured error code and message. * diff --git a/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php b/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php index 7513df3ff0fd5..b6a4a8115ae06 100644 --- a/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php +++ b/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php @@ -223,6 +223,84 @@ public function streamGenerateTextResult( array $prompt ): Generator { // phpcs: }; } + /** + * Creates a mock text generation model that returns scripted consecutive results. + * + * Each generateTextResult() call returns the next result from the list and + * records the message list it received in the referenced capture array. + * When the results run out, the last result is returned again. + * + * @param GenerativeAiResult[] $results The results to return, in order. + * @param array $captured_prompts Reference that receives the message list of each call. + * @param ModelMetadata|null $metadata Optional metadata. + * @return ModelInterface&TextGenerationModelInterface The mock model. + */ + protected function create_scripted_text_generation_model( + array $results, + array &$captured_prompts, + ?ModelMetadata $metadata = null + ): ModelInterface { + $metadata = $metadata ?? $this->create_test_text_model_metadata(); + + $provider_metadata = new ProviderMetadata( + 'mock', + 'Mock Provider', + ProviderTypeEnum::cloud() + ); + + return new class( $metadata, $provider_metadata, $results, $captured_prompts ) implements ModelInterface, TextGenerationModelInterface { + + private ModelMetadata $metadata; + private ProviderMetadata $provider_metadata; + private array $results; + private array $captured_prompts; + private ModelConfig $config; + + public function __construct( + ModelMetadata $metadata, + ProviderMetadata $provider_metadata, + array $results, + array &$captured_prompts + ) { + $this->metadata = $metadata; + $this->provider_metadata = $provider_metadata; + $this->results = $results; + $this->captured_prompts = &$captured_prompts; + $this->config = new ModelConfig(); + } + + public function metadata(): ModelMetadata { + return $this->metadata; + } + + public function providerMetadata(): ProviderMetadata { + return $this->provider_metadata; + } + + public function setConfig( ModelConfig $config ): void { + $this->config = $config; + } + + public function getConfig(): ModelConfig { + return $this->config; + } + + public function generateTextResult( array $prompt ): GenerativeAiResult { + $this->captured_prompts[] = $prompt; + + if ( count( $this->results ) > 1 ) { + return array_shift( $this->results ); + } + + return $this->results[0]; + } + + public function streamGenerateTextResult( array $prompt ): Generator { + yield $this->generateTextResult( $prompt ); + } + }; + } + /** * Creates a mock image generation model using anonymous class. * diff --git a/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php b/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php index 92eadebf95963..d81fbeae3348f 100644 --- a/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php +++ b/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php @@ -880,4 +880,121 @@ public function test_constructor_accepts_mixed_ability_types() { $response2 = $resolver->execute_ability( $call2 ); $this->assertArrayHasKey( 'success', $response2->getResponse() ); } + + /** + * Test that the pre-resolve filter can short-circuit the ability execution. + * + * @ticket 64865 + */ + public function test_pre_resolve_filter_short_circuits_execution() { + $injected = new FunctionResponse( + 'test-id', + 'wpab__wpaiclienttests__returns-error', + array( 'mocked' => true ) + ); + + add_filter( + 'wp_ai_client_pre_resolve_ability_call', + static function ( $pre ) use ( $injected ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter + return $injected; + } + ); + + $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/returns-error' ); + $call = new FunctionCall( + 'test-id', + 'wpab__wpaiclienttests__returns-error', + array() + ); + + $response = $resolver->execute_ability( $call ); + + $this->assertSame( $injected, $response, 'The injected response should be returned instead of executing the ability.' ); + } + + /** + * Test that a WP_Error from the pre-resolve filter becomes an error response. + * + * @ticket 64865 + */ + public function test_pre_resolve_filter_error_becomes_error_response() { + add_filter( + 'wp_ai_client_pre_resolve_ability_call', + static function ( $pre ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter + return new WP_Error( 'vetoed', 'This call is not allowed.' ); + } + ); + + $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/simple' ); + $call = new FunctionCall( + 'test-id', + 'wpab__wpaiclienttests__simple', + array() + ); + + $response = $resolver->execute_ability( $call ); + + $data = $response->getResponse(); + $this->assertSame( 'vetoed', $data['code'] ); + $this->assertSame( 'This call is not allowed.', $data['error'] ); + } + + /** + * Test that the resolved action fires with the call and the response. + * + * @ticket 64865 + */ + public function test_resolved_action_fires_with_call_and_response() { + $captured = array(); + add_action( + 'wp_ai_client_ability_call_resolved', + static function ( $call, $response ) use ( &$captured ) { + $captured[] = array( $call, $response ); + }, + 10, + 2 + ); + + $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/simple' ); + $call = new FunctionCall( + 'test-id', + 'wpab__wpaiclienttests__simple', + array() + ); + + $response = $resolver->execute_ability( $call ); + + $this->assertCount( 1, $captured ); + $this->assertSame( $call, $captured[0][0] ); + $this->assertSame( $response, $captured[0][1] ); + } + + /** + * Test that the resolved action also fires for error responses. + * + * @ticket 64865 + */ + public function test_resolved_action_fires_for_error_responses() { + $captured = array(); + add_action( + 'wp_ai_client_ability_call_resolved', + static function ( $call, $response ) use ( &$captured ) { + $captured[] = $response; + }, + 10, + 2 + ); + + $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/simple' ); + $call = new FunctionCall( + 'test-id', + 'wpab__wpaiclienttests__with-params', + array( 'title' => 'Test' ) + ); + + $response = $resolver->execute_ability( $call ); + + $this->assertCount( 1, $captured ); + $this->assertSame( 'ability_not_allowed', $captured[0]->getResponse()['code'] ); + } } diff --git a/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php new file mode 100644 index 0000000000000..c6f18d589f173 --- /dev/null +++ b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php @@ -0,0 +1,731 @@ +registry = $this->createMock( ProviderRegistry::class ); + } + + /** + * Creates a result whose message consists of the given function calls. + * + * @param array $calls List of id, function name, and arguments triples. + * @return GenerativeAiResult The result containing the function calls. + */ + private function create_function_call_result( array $calls ): GenerativeAiResult { + $parts = array(); + foreach ( $calls as $call ) { + $parts[] = new MessagePart( new FunctionCall( $call[0], $call[1], $call[2] ) ); + } + + $candidate = new Candidate( + new ModelMessage( $parts ), + FinishReasonEnum::toolCalls() + ); + + return new GenerativeAiResult( + 'function-call-result', + array( $candidate ), + new TokenUsage( 5, 7, 12 ), + new ProviderMetadata( 'mock', 'Mock Provider', ProviderTypeEnum::cloud() ), + $this->create_test_text_model_metadata() + ); + } + + /** + * Creates a prompt builder backed by a scripted model with resolution enabled. + * + * @param GenerativeAiResult[] $results The results the model returns, in order. + * @param array $captured_prompts Reference that receives the message list of each model call. + * @param string ...$abilities Ability names to register on the builder. + * @return WP_AI_Client_Prompt_Builder The prompt builder. + */ + private function create_resolution_builder( array $results, array &$captured_prompts, string ...$abilities ): WP_AI_Client_Prompt_Builder { + $model = $this->create_scripted_text_generation_model( $results, $captured_prompts ); + + $builder = new WP_AI_Client_Prompt_Builder( $this->registry, 'Test prompt' ); + $builder->using_model( $model ); + + if ( ! empty( $abilities ) ) { + $builder->using_abilities( ...$abilities ); + } + + return $builder; + } + + /** + * Returns the function name for a test ability. + * + * @param string $ability_name The ability name. + * @return string The function name exposed to the model. + */ + private function function_name( string $ability_name ): string { + return WP_AI_Client_Ability_Function_Resolver::ability_name_to_function_name( $ability_name ); + } + + /** + * Test that using_ability_resolution() is chainable. + * + * @ticket 64865 + */ + public function test_using_ability_resolution_is_chainable() { + $builder = new WP_AI_Client_Prompt_Builder( $this->registry, 'Test prompt' ); + + $this->assertSame( $builder, $builder->using_ability_resolution() ); + } + + /** + * Test that an invalid max_iterations option is rejected. + * + * @ticket 64865 + * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::using_ability_resolution + */ + public function test_using_ability_resolution_rejects_invalid_max_iterations() { + $builder = new WP_AI_Client_Prompt_Builder( $this->registry, 'Test prompt' ); + + $this->assertSame( $builder, $builder->using_ability_resolution( array( 'max_iterations' => 0 ) ) ); + } + + /** + * Test that an invalid option falls back to the filtered default. + * + * @ticket 64865 + * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::using_ability_resolution + */ + public function test_invalid_max_iterations_falls_back_to_filtered_default() { + add_filter( + 'wp_ai_client_ability_resolution_defaults', + static function ( $defaults ) { + $defaults['max_iterations'] = 2; + return $defaults; + } + ); + + $captured = array(); + $call_result = $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ); + + // The scripted model keeps returning the function call result. + $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' ); + $result = $builder + ->using_ability_resolution( array( 'max_iterations' => 0 ) ) + ->generate_text_result(); + + $this->assertSame( 2, $result->getAdditionalData()['ability_resolution']['rounds'] ); + } + + /** + * Test that the defaults filter can change the maximum number of rounds. + * + * @ticket 64865 + */ + public function test_resolution_defaults_can_be_filtered() { + add_filter( + 'wp_ai_client_ability_resolution_defaults', + static function ( $defaults ) { + $defaults['max_iterations'] = 1; + return $defaults; + } + ); + + $captured = array(); + $call_result = $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ); + + // The scripted model keeps returning the function call result. + $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' ); + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertInstanceOf( GenerativeAiResult::class, $result ); + $this->assertCount( 2, $captured, 'The model should be called once initially and once for the single allowed round.' ); + + $resolution = $result->getAdditionalData()['ability_resolution']; + $this->assertSame( 'max_iterations', $resolution['stop_reason'] ); + $this->assertSame( 1, $resolution['rounds'] ); + } + + /** + * Test that a response without function calls passes through with loop metadata. + * + * @ticket 64865 + */ + public function test_result_without_function_calls_passes_through_with_metadata() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( $this->create_test_result( 'Plain answer' ) ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertInstanceOf( GenerativeAiResult::class, $result ); + $this->assertSame( 'Plain answer', $result->toText() ); + $this->assertCount( 1, $captured ); + + $resolution = $result->getAdditionalData()['ability_resolution']; + $this->assertSame( 'completed', $resolution['stop_reason'] ); + $this->assertSame( 0, $resolution['rounds'] ); + $this->assertSame( array(), $resolution['resolved_calls'] ); + $this->assertCount( 2, $resolution['messages'], 'The transcript should contain the prompt and the final response.' ); + } + + /** + * Test that an ability call is executed and the final answer is returned. + * + * @ticket 64865 + */ + public function test_resolves_ability_call_and_returns_final_answer() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Final answer' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertInstanceOf( GenerativeAiResult::class, $result ); + $this->assertSame( 'Final answer', $result->toText() ); + $this->assertCount( 2, $captured ); + + $resolution = $result->getAdditionalData()['ability_resolution']; + $this->assertSame( 'completed', $resolution['stop_reason'] ); + $this->assertSame( 1, $resolution['rounds'] ); + $this->assertSame( + array( + array( + 'id' => 'call-1', + 'ability' => 'wpaiclienttests/simple', + ), + ), + $resolution['resolved_calls'] + ); + $this->assertCount( 4, $resolution['messages'], 'The transcript should contain the prompt, the call, the response, and the final answer.' ); + } + + /** + * Test that the follow-up request contains the expected conversation. + * + * @ticket 64865 + */ + public function test_second_request_contains_expected_transcript() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Final answer' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $builder->using_ability_resolution()->generate_text_result(); + + $messages = $captured[1]; + $this->assertCount( 3, $messages ); + $this->assertTrue( $messages[0]->getRole()->isUser(), 'The first message should be the user prompt.' ); + $this->assertTrue( $messages[1]->getRole()->isModel(), 'The second message should be the model response.' ); + $this->assertTrue( $messages[2]->getRole()->isUser(), 'The third message should carry the function responses.' ); + + $parts = $messages[2]->getParts(); + $this->assertCount( 1, $parts ); + + $response = $parts[0]->getFunctionResponse(); + $this->assertInstanceOf( FunctionResponse::class, $response ); + $this->assertSame( 'call-1', $response->getId() ); + $this->assertSame( $this->function_name( 'wpaiclienttests/simple' ), $response->getName() ); + $this->assertSame( array( 'success' => true ), $response->getResponse() ); + } + + /** + * Test that ability arguments from the model reach the ability. + * + * @ticket 64865 + */ + public function test_resolves_ability_call_with_arguments() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/with-params' ), array( 'title' => 'Hello' ) ) ) + ), + $this->create_test_result( 'Done' ), + ), + $captured, + 'wpaiclienttests/with-params' + ); + + $builder->using_ability_resolution()->generate_text_result(); + + $response = $captured[1][2]->getParts()[0]->getFunctionResponse(); + $this->assertSame( + array( + 'success' => true, + 'title' => 'Hello', + ), + $response->getResponse() + ); + } + + /** + * Test that all calls from one response are answered in a single message. + * + * @ticket 64865 + */ + public function test_answers_all_calls_from_one_response() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( + array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ), + array( 'call-2', $this->function_name( 'wpaiclienttests/with-params' ), array( 'title' => 'Hello' ) ), + ) + ), + $this->create_test_result( 'Done' ), + ), + $captured, + 'wpaiclienttests/simple', + 'wpaiclienttests/with-params' + ); + + $builder->using_ability_resolution()->generate_text_result(); + + $parts = $captured[1][2]->getParts(); + $this->assertCount( 2, $parts ); + $this->assertSame( 'call-1', $parts[0]->getFunctionResponse()->getId() ); + $this->assertSame( 'call-2', $parts[1]->getFunctionResponse()->getId() ); + } + + /** + * Test that the loop stops after the configured number of rounds. + * + * @ticket 64865 + */ + public function test_stops_after_max_iterations() { + $captured = array(); + $call_result = $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ); + + // The scripted model keeps returning the function call result. + $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' ); + $result = $builder + ->using_ability_resolution( array( 'max_iterations' => 2 ) ) + ->generate_text_result(); + + $this->assertInstanceOf( GenerativeAiResult::class, $result ); + $this->assertCount( 3, $captured, 'The model should be called once initially and once per allowed round.' ); + + $resolution = $result->getAdditionalData()['ability_resolution']; + $this->assertSame( 'max_iterations', $resolution['stop_reason'] ); + $this->assertSame( 2, $resolution['rounds'] ); + } + + /** + * Test that generate_text() returns an error when the loop is incomplete. + * + * @ticket 64865 + */ + public function test_generate_text_returns_error_when_max_iterations_reached() { + $captured = array(); + $call_result = $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ); + + $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' ); + $result = $builder + ->using_ability_resolution( array( 'max_iterations' => 1 ) ) + ->generate_text(); + + $this->assertWPError( $result ); + $this->assertSame( 'ability_resolution_incomplete', $result->get_error_code() ); + $this->assertSame( 'max_iterations', $result->get_error_data()['stop_reason'] ); + } + + /** + * Test that generate_text() returns the final answer through the loop. + * + * @ticket 64865 + */ + public function test_generate_text_returns_final_answer() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Final answer' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text(); + + $this->assertSame( 'Final answer', $result ); + } + + /** + * Test that the loop stops without executing anything when unknown functions are requested. + * + * @ticket 64865 + */ + public function test_stops_when_response_contains_unknown_function_calls() { + $resolved_fired = 0; + add_action( + 'wp_ai_client_ability_call_resolved', + static function () use ( &$resolved_fired ) { + ++$resolved_fired; + } + ); + + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( + array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ), + array( 'call-2', 'custom_function', array() ), + ) + ), + $this->create_test_result( 'Never returned' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertInstanceOf( GenerativeAiResult::class, $result ); + $this->assertCount( 1, $captured, 'The loop should not request a follow-up response.' ); + $this->assertSame( 0, $resolved_fired, 'No ability should be executed when unknown functions are requested.' ); + + $resolution = $result->getAdditionalData()['ability_resolution']; + $this->assertSame( 'unresolved_function_calls', $resolution['stop_reason'] ); + $this->assertSame( 0, $resolution['rounds'] ); + } + + /** + * Test that an ability error is sent back to the model and the loop continues. + * + * @ticket 64865 + */ + public function test_error_from_ability_is_sent_back_to_model() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/returns-error' ), array() ) ) + ), + $this->create_test_result( 'Recovered' ), + ), + $captured, + 'wpaiclienttests/returns-error' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( 'Recovered', $result->toText() ); + + $response = $captured[1][2]->getParts()[0]->getFunctionResponse()->getResponse(); + $this->assertSame( 'test_error', $response['code'] ); + } + + /** + * Test that a call to an ability outside the allowed list is answered with an error. + * + * @ticket 64865 + */ + public function test_not_allowed_ability_error_is_sent_back_to_model() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/with-params' ), array( 'title' => 'Hello' ) ) ) + ), + $this->create_test_result( 'Done' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( 'Done', $result->toText() ); + + $response = $captured[1][2]->getParts()[0]->getFunctionResponse()->getResponse(); + $this->assertSame( 'ability_not_allowed', $response['code'] ); + } + + /** + * Test that the prevent filter also stops the loop between rounds. + * + * @ticket 64865 + */ + public function test_prevent_filter_stops_the_loop_between_rounds() { + $evaluations = 0; + add_filter( + 'wp_ai_client_prevent_prompt', + static function ( $prevent ) use ( &$evaluations ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter + ++$evaluations; + return $evaluations > 1; + } + ); + + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Never returned' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertWPError( $result ); + $this->assertSame( 'prompt_prevented', $result->get_error_code() ); + $this->assertCount( 1, $captured, 'The follow-up request should be prevented.' ); + } + + /** + * Test that token usage is aggregated across all rounds. + * + * @ticket 64865 + */ + public function test_token_usage_is_aggregated_across_rounds() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + // Uses 5 prompt, 7 completion, and 12 total tokens. + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + // Uses 10 prompt, 20 completion, and 30 total tokens. + $this->create_test_result( 'Final answer' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + $usage = $result->getTokenUsage(); + + $this->assertSame( 15, $usage->getPromptTokens() ); + $this->assertSame( 27, $usage->getCompletionTokens() ); + $this->assertSame( 42, $usage->getTotalTokens() ); + } + + /** + * Test that lifecycle events fire for every round. + * + * @ticket 64865 + */ + public function test_lifecycle_events_fire_for_each_round() { + $before_fired = 0; + add_action( + 'wp_ai_client_before_generate_result', + static function ( $event ) use ( &$before_fired ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter + ++$before_fired; + } + ); + + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Final answer' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( 2, $before_fired, 'The before event should fire for the initial request and each round.' ); + } + + /** + * Test that ability declarations added directly are resolvable too. + * + * The allow-list is derived from the function declarations exposed to the + * model, so declarations built without using_abilities() participate. + * + * @ticket 64865 + */ + public function test_directly_declared_abilities_are_resolvable() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Final answer' ), + ), + $captured + ); + + $builder->using_function_declarations( + new FunctionDeclaration( $this->function_name( 'wpaiclienttests/simple' ), 'A simple test ability.' ) + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( 'Final answer', $result->toText() ); + $this->assertSame( 1, $result->getAdditionalData()['ability_resolution']['rounds'] ); + } + + /** + * Test that replacing the declarations also replaces the allow-list. + * + * @ticket 64865 + */ + public function test_replaced_declarations_limit_the_allow_list() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + $this->create_test_result( 'Done' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + // Replaces the declarations from using_abilities() above. + $builder->using_function_declarations( + new FunctionDeclaration( $this->function_name( 'wpaiclienttests/with-params' ), 'Another test ability.' ) + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( 'Done', $result->toText() ); + + // The call to the no longer declared ability must not execute. + $response = $captured[1][2]->getParts()[0]->getFunctionResponse()->getResponse(); + $this->assertSame( 'ability_not_allowed', $response['code'] ); + } + + /** + * Test that resolution without registered abilities falls back to plain generation. + * + * @ticket 64865 + * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::generate_with_ability_resolution + */ + public function test_resolution_without_abilities_falls_back_to_plain_generation() { + $captured = array(); + $text_result = $this->create_test_result( 'Plain answer' ); + + $builder = $this->create_resolution_builder( array( $text_result ), $captured ); + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( $text_result, $result, 'The unmodified result should be returned.' ); + $this->assertCount( 1, $captured ); + } + + /** + * Test that the generate_text() fallback still returns a string. + * + * @ticket 64865 + * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::generate_with_ability_resolution + */ + public function test_generate_text_without_abilities_falls_back_to_plain_text() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( $this->create_test_result( 'Plain answer' ) ), + $captured + ); + + $result = $builder->using_ability_resolution()->generate_text(); + + $this->assertSame( 'Plain answer', $result ); + } + + /** + * Test that unsupported generation methods warn and run without resolution. + * + * @ticket 64865 + * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::__call + */ + public function test_resolution_warns_for_unsupported_generation_methods() { + $captured = array(); + $builder = $this->create_resolution_builder( + array( $this->create_test_result( 'Plain answer' ) ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_image_result(); + + $this->assertWPError( $result ); + } +} From cf4f2e6170311634e44921499b61dc9f3f3e348f Mon Sep 17 00:00:00 2001 From: Grzegorz Ziolkowski Date: Thu, 23 Jul 2026 14:44:11 +0200 Subject: [PATCH 2/3] AI: Remove ability resolution hooks. --- ...wp-ai-client-ability-function-resolver.php | 58 --------- .../class-wp-ai-client-prompt-builder.php | 29 +---- .../wpAiClientAbilityFunctionResolver.php | 117 ------------------ .../ai-client/wpAiClientAbilityResolution.php | 55 ++------ 4 files changed, 14 insertions(+), 245 deletions(-) diff --git a/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php b/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php index 0a755539d4340..9f3d8909afbfe 100644 --- a/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php +++ b/src/wp-includes/ai-client/class-wp-ai-client-ability-function-resolver.php @@ -103,34 +103,6 @@ public static function is_ability_function_name( string $function_name ): bool { * @return FunctionResponse The response from executing the ability. */ public function execute_ability( FunctionCall $call ): FunctionResponse { - $response = $this->resolve_ability_call( $call ); - - /** - * Fires after an ability function call has been resolved. - * - * The response may represent a successful execution or an error, for - * example when the ability is not allowed or its execution failed. - * Useful for logging and auditing ability usage by AI models. - * - * @since 7.2.0 - * - * @param FunctionCall $call The function call requested by the model. - * @param FunctionResponse $response The response that will be sent back to the model. - */ - do_action( 'wp_ai_client_ability_call_resolved', $call, $response ); - - return $response; - } - - /** - * Resolves a function call into a function response. - * - * @since 7.2.0 - * - * @param FunctionCall $call The function call to resolve. - * @return FunctionResponse The response from resolving the call. - */ - private function resolve_ability_call( FunctionCall $call ): FunctionResponse { $function_name = $call->getName() ?? 'unknown'; $function_id = $call->getId() ?? 'unknown'; @@ -173,36 +145,6 @@ private function resolve_ability_call( FunctionCall $call ): FunctionResponse { ); } - /** - * Filters whether to short-circuit the execution of an ability function call. - * - * Returning a FunctionResponse skips the ability execution and sends that - * response back to the model instead. Returning a WP_Error skips the - * execution and sends an error response back to the model. This allows - * vetoing individual calls, for example based on ability annotations. - * - * @since 7.2.0 - * - * @param FunctionResponse|WP_Error|null $pre A response to short-circuit with, or null to execute the ability. Default null. - * @param FunctionCall $call The function call requested by the model. - * @param WP_Ability $ability The ability that is about to be executed. - */ - $pre = apply_filters( 'wp_ai_client_pre_resolve_ability_call', null, $call, $ability ); - if ( $pre instanceof FunctionResponse ) { - return $pre; - } - if ( is_wp_error( $pre ) ) { - return new FunctionResponse( - $function_id, - $function_name, - array( - 'error' => $pre->get_error_message(), - 'code' => $pre->get_error_code(), - 'data' => $pre->get_error_data(), - ) - ); - } - $args = $call->getArgs(); $result = $ability->execute( ! empty( $args ) ? $args : null ); diff --git a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php index 84024e3c35e82..8287a1403bd72 100644 --- a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php +++ b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php @@ -319,30 +319,13 @@ public function using_abilities( ...$abilities ): self { * @return self The current instance for method chaining. */ public function using_ability_resolution( array $options = array() ): self { - $defaults = array( - 'max_iterations' => 5, + $options = wp_parse_args( + $options, + array( + 'max_iterations' => 5, + ) ); - /** - * Filters the default options for automatic ability resolution. - * - * @since 7.2.0 - * - * @param array $defaults Default options. See {@see WP_AI_Client_Prompt_Builder::using_ability_resolution()} - * for the supported keys. - */ - $filtered_defaults = apply_filters( 'wp_ai_client_ability_resolution_defaults', $defaults ); - if ( is_array( $filtered_defaults ) ) { - $defaults = wp_parse_args( $filtered_defaults, $defaults ); - } - - // Guard against invalid filtered defaults. - if ( ! is_int( $defaults['max_iterations'] ) || $defaults['max_iterations'] < 1 ) { - $defaults['max_iterations'] = 5; - } - - $options = wp_parse_args( $options, $defaults ); - if ( ! is_int( $options['max_iterations'] ) || $options['max_iterations'] < 1 ) { _doing_it_wrong( __METHOD__, @@ -353,7 +336,7 @@ public function using_ability_resolution( array $options = array() ): self { ), '7.2.0' ); - $options['max_iterations'] = $defaults['max_iterations']; + $options['max_iterations'] = 5; } $this->ability_resolution_options = array( diff --git a/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php b/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php index d81fbeae3348f..92eadebf95963 100644 --- a/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php +++ b/tests/phpunit/tests/ai-client/wpAiClientAbilityFunctionResolver.php @@ -880,121 +880,4 @@ public function test_constructor_accepts_mixed_ability_types() { $response2 = $resolver->execute_ability( $call2 ); $this->assertArrayHasKey( 'success', $response2->getResponse() ); } - - /** - * Test that the pre-resolve filter can short-circuit the ability execution. - * - * @ticket 64865 - */ - public function test_pre_resolve_filter_short_circuits_execution() { - $injected = new FunctionResponse( - 'test-id', - 'wpab__wpaiclienttests__returns-error', - array( 'mocked' => true ) - ); - - add_filter( - 'wp_ai_client_pre_resolve_ability_call', - static function ( $pre ) use ( $injected ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter - return $injected; - } - ); - - $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/returns-error' ); - $call = new FunctionCall( - 'test-id', - 'wpab__wpaiclienttests__returns-error', - array() - ); - - $response = $resolver->execute_ability( $call ); - - $this->assertSame( $injected, $response, 'The injected response should be returned instead of executing the ability.' ); - } - - /** - * Test that a WP_Error from the pre-resolve filter becomes an error response. - * - * @ticket 64865 - */ - public function test_pre_resolve_filter_error_becomes_error_response() { - add_filter( - 'wp_ai_client_pre_resolve_ability_call', - static function ( $pre ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter - return new WP_Error( 'vetoed', 'This call is not allowed.' ); - } - ); - - $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/simple' ); - $call = new FunctionCall( - 'test-id', - 'wpab__wpaiclienttests__simple', - array() - ); - - $response = $resolver->execute_ability( $call ); - - $data = $response->getResponse(); - $this->assertSame( 'vetoed', $data['code'] ); - $this->assertSame( 'This call is not allowed.', $data['error'] ); - } - - /** - * Test that the resolved action fires with the call and the response. - * - * @ticket 64865 - */ - public function test_resolved_action_fires_with_call_and_response() { - $captured = array(); - add_action( - 'wp_ai_client_ability_call_resolved', - static function ( $call, $response ) use ( &$captured ) { - $captured[] = array( $call, $response ); - }, - 10, - 2 - ); - - $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/simple' ); - $call = new FunctionCall( - 'test-id', - 'wpab__wpaiclienttests__simple', - array() - ); - - $response = $resolver->execute_ability( $call ); - - $this->assertCount( 1, $captured ); - $this->assertSame( $call, $captured[0][0] ); - $this->assertSame( $response, $captured[0][1] ); - } - - /** - * Test that the resolved action also fires for error responses. - * - * @ticket 64865 - */ - public function test_resolved_action_fires_for_error_responses() { - $captured = array(); - add_action( - 'wp_ai_client_ability_call_resolved', - static function ( $call, $response ) use ( &$captured ) { - $captured[] = $response; - }, - 10, - 2 - ); - - $resolver = new WP_AI_Client_Ability_Function_Resolver( 'wpaiclienttests/simple' ); - $call = new FunctionCall( - 'test-id', - 'wpab__wpaiclienttests__with-params', - array( 'title' => 'Test' ) - ); - - $response = $resolver->execute_ability( $call ); - - $this->assertCount( 1, $captured ); - $this->assertSame( 'ability_not_allowed', $captured[0]->getResponse()['code'] ); - } } diff --git a/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php index c6f18d589f173..1f55e4ba355bb 100644 --- a/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php +++ b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php @@ -139,20 +139,12 @@ public function test_using_ability_resolution_rejects_invalid_max_iterations() { } /** - * Test that an invalid option falls back to the filtered default. + * Test that an invalid option falls back to the default. * * @ticket 64865 * @expectedIncorrectUsage WP_AI_Client_Prompt_Builder::using_ability_resolution */ - public function test_invalid_max_iterations_falls_back_to_filtered_default() { - add_filter( - 'wp_ai_client_ability_resolution_defaults', - static function ( $defaults ) { - $defaults['max_iterations'] = 2; - return $defaults; - } - ); - + public function test_invalid_max_iterations_falls_back_to_default() { $captured = array(); $call_result = $this->create_function_call_result( array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) @@ -164,38 +156,7 @@ static function ( $defaults ) { ->using_ability_resolution( array( 'max_iterations' => 0 ) ) ->generate_text_result(); - $this->assertSame( 2, $result->getAdditionalData()['ability_resolution']['rounds'] ); - } - - /** - * Test that the defaults filter can change the maximum number of rounds. - * - * @ticket 64865 - */ - public function test_resolution_defaults_can_be_filtered() { - add_filter( - 'wp_ai_client_ability_resolution_defaults', - static function ( $defaults ) { - $defaults['max_iterations'] = 1; - return $defaults; - } - ); - - $captured = array(); - $call_result = $this->create_function_call_result( - array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) - ); - - // The scripted model keeps returning the function call result. - $builder = $this->create_resolution_builder( array( $call_result ), $captured, 'wpaiclienttests/simple' ); - $result = $builder->using_ability_resolution()->generate_text_result(); - - $this->assertInstanceOf( GenerativeAiResult::class, $result ); - $this->assertCount( 2, $captured, 'The model should be called once initially and once for the single allowed round.' ); - - $resolution = $result->getAdditionalData()['ability_resolution']; - $this->assertSame( 'max_iterations', $resolution['stop_reason'] ); - $this->assertSame( 1, $resolution['rounds'] ); + $this->assertSame( 5, $result->getAdditionalData()['ability_resolution']['rounds'] ); } /** @@ -434,11 +395,11 @@ public function test_generate_text_returns_final_answer() { * @ticket 64865 */ public function test_stops_when_response_contains_unknown_function_calls() { - $resolved_fired = 0; + $invoked_abilities = array(); add_action( - 'wp_ai_client_ability_call_resolved', - static function () use ( &$resolved_fired ) { - ++$resolved_fired; + 'wp_ability_invoked', + static function ( $ability_name ) use ( &$invoked_abilities ) { + $invoked_abilities[] = $ability_name; } ); @@ -461,7 +422,7 @@ static function () use ( &$resolved_fired ) { $this->assertInstanceOf( GenerativeAiResult::class, $result ); $this->assertCount( 1, $captured, 'The loop should not request a follow-up response.' ); - $this->assertSame( 0, $resolved_fired, 'No ability should be executed when unknown functions are requested.' ); + $this->assertSame( array(), $invoked_abilities, 'No ability should be executed when unknown functions are requested.' ); $resolution = $result->getAdditionalData()['ability_resolution']; $this->assertSame( 'unresolved_function_calls', $resolution['stop_reason'] ); From 15a18ec6e1a03cbfc8e0dcf0cf392b86ad03d3fb Mon Sep 17 00:00:00 2001 From: Grzegorz Ziolkowski Date: Thu, 23 Jul 2026 15:18:11 +0200 Subject: [PATCH 3/3] AI: Harden ability resolution loop behavior and tests. --- .../class-wp-ai-client-prompt-builder.php | 12 +- ...wp-ai-client-mock-model-creation-trait.php | 21 +++- .../ai-client/wpAiClientAbilityResolution.php | 110 ++++++++++++++++-- 3 files changed, 123 insertions(+), 20 deletions(-) diff --git a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php index 8287a1403bd72..a683b54755cbc 100644 --- a/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php +++ b/src/wp-includes/ai-client/class-wp-ai-client-prompt-builder.php @@ -592,6 +592,12 @@ private function generate_with_ability_resolution( string $method ) { break; } + $prevented = $this->get_prompt_prevented_error(); + if ( null !== $prevented ) { + $this->error = $prevented; + return $this->error; + } + $responses = $resolver->execute_abilities( $message ); foreach ( $ability_calls as $call ) { @@ -605,12 +611,6 @@ private function generate_with_ability_resolution( string $method ) { $transcript[] = $responses; ++$rounds; - $prevented = $this->get_prompt_prevented_error(); - if ( null !== $prevented ) { - $this->error = $prevented; - return $this->error; - } - if ( null !== $dispatcher ) { $dispatcher->dispatch( new BeforeGenerateResultEvent( $transcript, $model, $capability ) ); } diff --git a/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php b/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php index b6a4a8115ae06..3ec9b26d3f55b 100644 --- a/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php +++ b/tests/phpunit/includes/wp-ai-client-mock-model-creation-trait.php @@ -230,16 +230,21 @@ public function streamGenerateTextResult( array $prompt ): Generator { // phpcs: * records the message list it received in the referenced capture array. * When the results run out, the last result is returned again. * - * @param GenerativeAiResult[] $results The results to return, in order. - * @param array $captured_prompts Reference that receives the message list of each call. - * @param ModelMetadata|null $metadata Optional metadata. + * @param array $results Scripted results or exceptions, in order. + * @param array $captured_prompts Receives the message list of each call. + * @param ModelMetadata|null $metadata Optional metadata. * @return ModelInterface&TextGenerationModelInterface The mock model. + * @throws InvalidArgumentException If no results are provided. */ protected function create_scripted_text_generation_model( array $results, array &$captured_prompts, ?ModelMetadata $metadata = null ): ModelInterface { + if ( empty( $results ) ) { + throw new InvalidArgumentException( 'At least one scripted result is required.' ); + } + $metadata = $metadata ?? $this->create_test_text_model_metadata(); $provider_metadata = new ProviderMetadata( @@ -289,10 +294,16 @@ public function generateTextResult( array $prompt ): GenerativeAiResult { $this->captured_prompts[] = $prompt; if ( count( $this->results ) > 1 ) { - return array_shift( $this->results ); + $result = array_shift( $this->results ); + } else { + $result = $this->results[0]; + } + + if ( $result instanceof Exception ) { + throw $result; } - return $this->results[0]; + return $result; } public function streamGenerateTextResult( array $prompt ): Generator { diff --git a/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php index 1f55e4ba355bb..42b11ad6ecfeb 100644 --- a/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php +++ b/tests/phpunit/tests/ai-client/wpAiClientAbilityResolution.php @@ -6,6 +6,8 @@ * @covers WP_AI_Client_Prompt_Builder */ +use WordPress\AiClient\Events\AfterGenerateResultEvent; +use WordPress\AiClient\Events\BeforeGenerateResultEvent; use WordPress\AiClient\Messages\DTO\MessagePart; use WordPress\AiClient\Messages\DTO\ModelMessage; use WordPress\AiClient\Providers\DTO\ProviderMetadata; @@ -87,9 +89,9 @@ private function create_function_call_result( array $calls ): GenerativeAiResult /** * Creates a prompt builder backed by a scripted model with resolution enabled. * - * @param GenerativeAiResult[] $results The results the model returns, in order. - * @param array $captured_prompts Reference that receives the message list of each model call. - * @param string ...$abilities Ability names to register on the builder. + * @param array $results Scripted results or exceptions, in order. + * @param array $captured_prompts Receives each model call's message list. + * @param string ...$abilities Ability names to register on the builder. * @return WP_AI_Client_Prompt_Builder The prompt builder. */ private function create_resolution_builder( array $results, array &$captured_prompts, string ...$abilities ): WP_AI_Client_Prompt_Builder { @@ -126,6 +128,20 @@ public function test_using_ability_resolution_is_chainable() { $this->assertSame( $builder, $builder->using_ability_resolution() ); } + /** + * Test that the scripted model requires at least one result. + * + * @ticket 64865 + */ + public function test_scripted_model_requires_at_least_one_result() { + $captured = array(); + + $this->expectException( InvalidArgumentException::class ); + $this->expectExceptionMessage( 'At least one scripted result is required.' ); + + $this->create_scripted_text_generation_model( array(), $captured ); + } + /** * Test that an invalid max_iterations option is rejected. * @@ -487,7 +503,8 @@ public function test_not_allowed_ability_error_is_sent_back_to_model() { * @ticket 64865 */ public function test_prevent_filter_stops_the_loop_between_rounds() { - $evaluations = 0; + $evaluations = 0; + $invoked_abilities = array(); add_filter( 'wp_ai_client_prevent_prompt', static function ( $prevent ) use ( &$evaluations ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter @@ -495,6 +512,12 @@ static function ( $prevent ) use ( &$evaluations ) { // phpcs:ignore Generic.Cod return $evaluations > 1; } ); + add_action( + 'wp_ability_invoked', + static function ( $ability_name ) use ( &$invoked_abilities ) { + $invoked_abilities[] = $ability_name; + } + ); $captured = array(); $builder = $this->create_resolution_builder( @@ -513,6 +536,7 @@ static function ( $prevent ) use ( &$evaluations ) { // phpcs:ignore Generic.Cod $this->assertWPError( $result ); $this->assertSame( 'prompt_prevented', $result->get_error_code() ); $this->assertCount( 1, $captured, 'The follow-up request should be prevented.' ); + $this->assertSame( array(), $invoked_abilities, 'No ability should be executed after prompt execution is prevented.' ); } /** @@ -549,11 +573,18 @@ public function test_token_usage_is_aggregated_across_rounds() { * @ticket 64865 */ public function test_lifecycle_events_fire_for_each_round() { - $before_fired = 0; + $before_events = array(); + $after_events = array(); add_action( 'wp_ai_client_before_generate_result', - static function ( $event ) use ( &$before_fired ) { // phpcs:ignore Generic.CodeAnalysis.UnusedFunctionParameter - ++$before_fired; + static function ( $event ) use ( &$before_events ) { + $before_events[] = $event; + } + ); + add_action( + 'wp_ai_client_after_generate_result', + static function ( $event ) use ( &$after_events ) { + $after_events[] = $event; } ); @@ -569,9 +600,70 @@ static function ( $event ) use ( &$before_fired ) { // phpcs:ignore Generic.Code 'wpaiclienttests/simple' ); - $builder->using_ability_resolution()->generate_text_result(); + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertSame( 'Final answer', $result->toText() ); + $this->assertCount( 2, $before_events, 'The before event should fire for the initial request and each round.' ); + $this->assertCount( 2, $after_events, 'The after event should fire after the initial request and each successful round.' ); + + $this->assertInstanceOf( BeforeGenerateResultEvent::class, $before_events[0] ); + $this->assertInstanceOf( BeforeGenerateResultEvent::class, $before_events[1] ); + $this->assertCount( 1, $before_events[0]->getMessages() ); + $this->assertCount( 3, $before_events[1]->getMessages() ); + $this->assertEquals( $captured[0], $before_events[0]->getMessages() ); + $this->assertEquals( $captured[1], $before_events[1]->getMessages() ); + + $this->assertInstanceOf( AfterGenerateResultEvent::class, $after_events[0] ); + $this->assertInstanceOf( AfterGenerateResultEvent::class, $after_events[1] ); + $this->assertCount( 1, $after_events[0]->getMessages() ); + $this->assertCount( 3, $after_events[1]->getMessages() ); + $this->assertCount( 1, $after_events[0]->getResult()->toMessage()->getParts() ); + $this->assertSame( 'call-1', $after_events[0]->getResult()->toMessage()->getParts()[0]->getFunctionCall()->getId() ); + $this->assertSame( 'Final answer', $after_events[1]->getResult()->toText() ); + } - $this->assertSame( 2, $before_fired, 'The before event should fire for the initial request and each round.' ); + /** + * Test that a failed follow-up request does not dispatch an after event. + * + * @ticket 64865 + */ + public function test_failed_follow_up_request_does_not_fire_after_event() { + $before_events = array(); + $after_events = array(); + add_action( + 'wp_ai_client_before_generate_result', + static function ( $event ) use ( &$before_events ) { + $before_events[] = $event; + } + ); + add_action( + 'wp_ai_client_after_generate_result', + static function ( $event ) use ( &$after_events ) { + $after_events[] = $event; + } + ); + + $captured = array(); + $builder = $this->create_resolution_builder( + array( + $this->create_function_call_result( + array( array( 'call-1', $this->function_name( 'wpaiclienttests/simple' ), array() ) ) + ), + new RuntimeException( 'Follow-up failed.' ), + ), + $captured, + 'wpaiclienttests/simple' + ); + + $result = $builder->using_ability_resolution()->generate_text_result(); + + $this->assertWPError( $result ); + $this->assertSame( 'prompt_builder_error', $result->get_error_code() ); + $this->assertSame( 'Follow-up failed.', $result->get_error_message() ); + $this->assertCount( 2, $captured ); + $this->assertCount( 2, $before_events, 'The before event should fire before the failed follow-up request.' ); + $this->assertCount( 1, $after_events, 'The after event should only fire for the successful initial request.' ); + $this->assertInstanceOf( AfterGenerateResultEvent::class, $after_events[0] ); } /**