[Summary view] [Print] [Text view]

   1  <?php
   2  /**
   3   * WP AI Client: WP_AI_Client_Ability_Function_Resolver class
   4   *
   5   * @package WordPress
   6   * @subpackage AI
   7   * @since 7.0.0
   8   */
   9  
  10  use WordPress\AiClient\Messages\DTO\Message;
  11  use WordPress\AiClient\Messages\DTO\MessagePart;
  12  use WordPress\AiClient\Messages\DTO\UserMessage;
  13  use WordPress\AiClient\Tools\DTO\FunctionCall;
  14  use WordPress\AiClient\Tools\DTO\FunctionResponse;
  15  
  16  /**
  17   * Resolves and executes WordPress Abilities API function calls from AI models.
  18   *
  19   * This class must be instantiated with the specific abilities that the AI model
  20   * is allowed to execute, ensuring that only explicitly specified abilities can
  21   * be called. This prevents the model from executing arbitrary abilities.
  22   *
  23   * @since 7.0.0
  24   */
  25  class WP_AI_Client_Ability_Function_Resolver {
  26  
  27      /**
  28       * Prefix used to identify ability function calls.
  29       *
  30       * @since 7.0.0
  31       * @var string
  32       */
  33      private const ABILITY_PREFIX = 'wpab__';
  34  
  35      /**
  36       * Map of allowed ability names for this instance.
  37       *
  38       * Keys are ability name strings, values are `true` for O(1) lookup.
  39       *
  40       * @since 7.0.0
  41       * @var array<string, true>
  42       */
  43      private array $allowed_abilities;
  44  
  45      /**
  46       * Constructor.
  47       *
  48       * @since 7.0.0
  49       *
  50       * @param WP_Ability|string ...$abilities The abilities that this resolver is allowed to execute.
  51       */
  52  	public function __construct( ...$abilities ) {
  53          $this->allowed_abilities = array();
  54  
  55          foreach ( $abilities as $ability ) {
  56              if ( $ability instanceof WP_Ability ) {
  57                  $this->allowed_abilities[ $ability->get_name() ] = true;
  58              } elseif ( is_string( $ability ) ) {
  59                  $this->allowed_abilities[ $ability ] = true;
  60              }
  61          }
  62      }
  63  
  64      /**
  65       * Checks if a function call is an ability call.
  66       *
  67       * @since 7.0.0
  68       *
  69       * @param FunctionCall $call The function call to check.
  70       * @return bool True if the function call is an ability call, false otherwise.
  71       */
  72  	public function is_ability_call( FunctionCall $call ): bool {
  73          $name = $call->getName();
  74          if ( null === $name ) {
  75              return false;
  76          }
  77  
  78          return str_starts_with( $name, self::ABILITY_PREFIX );
  79      }
  80  
  81      /**
  82       * Executes a WordPress ability from a function call.
  83       *
  84       * Only abilities that were specified in the constructor are allowed to be
  85       * executed. If the ability is not in the allowed list, an error response
  86       * with code `ability_not_allowed` is returned.
  87       *
  88       * @since 7.0.0
  89       *
  90       * @param FunctionCall $call The function call to execute.
  91       * @return FunctionResponse The response from executing the ability.
  92       */
  93  	public function execute_ability( FunctionCall $call ): FunctionResponse {
  94          $function_name = $call->getName() ?? 'unknown';
  95          $function_id   = $call->getId() ?? 'unknown';
  96  
  97          if ( ! $this->is_ability_call( $call ) ) {
  98              return new FunctionResponse(
  99                  $function_id,
 100                  $function_name,
 101                  array(
 102                      'error' => __( 'Not an ability function call' ),
 103                      'code'  => 'invalid_ability_call',
 104                  )
 105              );
 106          }
 107  
 108          $ability_name = self::function_name_to_ability_name( $function_name );
 109  
 110          if ( ! isset( $this->allowed_abilities[ $ability_name ] ) ) {
 111              return new FunctionResponse(
 112                  $function_id,
 113                  $function_name,
 114                  array(
 115                      /* translators: %s: ability name */
 116                      'error' => sprintf( __( 'Ability "%s" was not specified in the allowed abilities list.' ), $ability_name ),
 117                      'code'  => 'ability_not_allowed',
 118                  )
 119              );
 120          }
 121  
 122          $ability = wp_get_ability( $ability_name );
 123  
 124          if ( ! $ability instanceof WP_Ability ) {
 125              return new FunctionResponse(
 126                  $function_id,
 127                  $function_name,
 128                  array(
 129                      /* translators: %s: ability name */
 130                      'error' => sprintf( __( 'Ability "%s" not found' ), $ability_name ),
 131                      'code'  => 'ability_not_found',
 132                  )
 133              );
 134          }
 135  
 136          $args   = $call->getArgs();
 137          $result = $ability->execute( ! empty( $args ) ? $args : null );
 138  
 139          if ( is_wp_error( $result ) ) {
 140              return new FunctionResponse(
 141                  $function_id,
 142                  $function_name,
 143                  array(
 144                      'error' => $result->get_error_message(),
 145                      'code'  => $result->get_error_code(),
 146                      'data'  => $result->get_error_data(),
 147                  )
 148              );
 149          }
 150  
 151          $result = wp_prepare_json_value_for_client( $result, $ability->get_output_schema() );
 152  
 153          return new FunctionResponse(
 154              $function_id,
 155              $function_name,
 156              $result
 157          );
 158      }
 159  
 160      /**
 161       * Checks if a message contains any ability function calls.
 162       *
 163       * @since 7.0.0
 164       *
 165       * @param Message $message The message to check.
 166       * @return bool True if the message contains ability calls, false otherwise.
 167       */
 168  	public function has_ability_calls( Message $message ): bool {
 169          foreach ( $message->getParts() as $part ) {
 170              if ( $part->getType()->isFunctionCall() ) {
 171                  $function_call = $part->getFunctionCall();
 172                  if ( $function_call instanceof FunctionCall && $this->is_ability_call( $function_call ) ) {
 173                      return true;
 174                  }
 175              }
 176          }
 177  
 178          return false;
 179      }
 180  
 181      /**
 182       * Executes all ability function calls in a message.
 183       *
 184       * @since 7.0.0
 185       *
 186       * @param Message $message The message containing function calls.
 187       * @return Message A new message with function responses.
 188       */
 189  	public function execute_abilities( Message $message ): Message {
 190          $response_parts = array();
 191  
 192          foreach ( $message->getParts() as $part ) {
 193              if ( $part->getType()->isFunctionCall() ) {
 194                  $function_call = $part->getFunctionCall();
 195                  if ( $function_call instanceof FunctionCall && $this->is_ability_call( $function_call ) ) {
 196                      $function_response = $this->execute_ability( $function_call );
 197                      $response_parts[]  = new MessagePart( $function_response );
 198                  }
 199              }
 200          }
 201  
 202          return new UserMessage( $response_parts );
 203      }
 204  
 205      /**
 206       * Converts an ability name to a function name.
 207       *
 208       * Transforms "tec/create_event" to "wpab__tec__create_event".
 209       *
 210       * @since 7.0.0
 211       *
 212       * @param string $ability_name The ability name to convert.
 213       * @return string The function name.
 214       */
 215  	public static function ability_name_to_function_name( string $ability_name ): string {
 216          return self::ABILITY_PREFIX . str_replace( '/', '__', $ability_name );
 217      }
 218  
 219      /**
 220       * Converts a function name to an ability name.
 221       *
 222       * Transforms "wpab__tec__create_event" to "tec/create_event".
 223       *
 224       * @since 7.0.0
 225       *
 226       * @param string $function_name The function name to convert.
 227       * @return string The ability name.
 228       */
 229  	public static function function_name_to_ability_name( string $function_name ): string {
 230          $without_prefix = substr( $function_name, strlen( self::ABILITY_PREFIX ) );
 231  
 232          return str_replace( '__', '/', $without_prefix );
 233      }
 234  }