[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          return new FunctionResponse(
 152              $function_id,
 153              $function_name,
 154              $result
 155          );
 156      }
 157  
 158      /**
 159       * Checks if a message contains any ability function calls.
 160       *
 161       * @since 7.0.0
 162       *
 163       * @param Message $message The message to check.
 164       * @return bool True if the message contains ability calls, false otherwise.
 165       */
 166  	public function has_ability_calls( Message $message ): bool {
 167          foreach ( $message->getParts() as $part ) {
 168              if ( $part->getType()->isFunctionCall() ) {
 169                  $function_call = $part->getFunctionCall();
 170                  if ( $function_call instanceof FunctionCall && $this->is_ability_call( $function_call ) ) {
 171                      return true;
 172                  }
 173              }
 174          }
 175  
 176          return false;
 177      }
 178  
 179      /**
 180       * Executes all ability function calls in a message.
 181       *
 182       * @since 7.0.0
 183       *
 184       * @param Message $message The message containing function calls.
 185       * @return Message A new message with function responses.
 186       */
 187  	public function execute_abilities( Message $message ): Message {
 188          $response_parts = array();
 189  
 190          foreach ( $message->getParts() as $part ) {
 191              if ( $part->getType()->isFunctionCall() ) {
 192                  $function_call = $part->getFunctionCall();
 193                  if ( $function_call instanceof FunctionCall ) {
 194                      $function_response = $this->execute_ability( $function_call );
 195                      $response_parts[]  = new MessagePart( $function_response );
 196                  }
 197              }
 198          }
 199  
 200          return new UserMessage( $response_parts );
 201      }
 202  
 203      /**
 204       * Converts an ability name to a function name.
 205       *
 206       * Transforms "tec/create_event" to "wpab__tec__create_event".
 207       *
 208       * @since 7.0.0
 209       *
 210       * @param string $ability_name The ability name to convert.
 211       * @return string The function name.
 212       */
 213  	public static function ability_name_to_function_name( string $ability_name ): string {
 214          return self::ABILITY_PREFIX . str_replace( '/', '__', $ability_name );
 215      }
 216  
 217      /**
 218       * Converts a function name to an ability name.
 219       *
 220       * Transforms "wpab__tec__create_event" to "tec/create_event".
 221       *
 222       * @since 7.0.0
 223       *
 224       * @param string $function_name The function name to convert.
 225       * @return string The ability name.
 226       */
 227  	public static function function_name_to_ability_name( string $function_name ): string {
 228          $without_prefix = substr( $function_name, strlen( self::ABILITY_PREFIX ) );
 229  
 230          return str_replace( '__', '/', $without_prefix );
 231      }
 232  }