Issue #3483312 by nterbogt, mortim07: Add facet support

README.md

@@ -12,3 +12,7 @@ This module requires composer for installation. To install, simply run `composer
 ## Usage
 
 Start searching your index. Some example queries are included in `examples/` directory of the module.
+
+## Limitations
+
+Currently only string based facets are supported.
\ No newline at end of file

examples/autocomplete_query.graphql

 {
-  vertexAiAutocompleteQuery(query:"party") {
+  vertexAiAutocompleteQuery(query: "party") {
     suggestions
   }
 }

examples/search_query.graphql

 {
-  vertexAiSearchQuery(query:"baby",pageSize:5,offset:10,filter:"") {
+  vertexAiSearchQuery(query: "party", pageSize: 5, offset: 10, filter: "myKey: ANY(\"123\")") {
     totalResults
     summary
     results {
@@ -7,5 +7,12 @@
       uri
       snippet
     }
+    facets (keys: "myKey, myKeyNew") {
+      key
+      values {
+        value
+        count
+      }
+    }
   }
 }

graphql/vertex_ai_search_query.base.graphqls

@@ -2,6 +2,9 @@ type VertexAiSearchResponse {
     totalResults: Int!
     summary: String
     results: [VertexAiSearchResult]
+    facets(
+        keys: String!
+    ): [VertexAiFacet]
 }
 
 type VertexAiSearchResult {
@@ -12,6 +15,16 @@ type VertexAiSearchResult {
     snippet: String
 }
 
+type VertexAiFacet {
+    key: String!
+    values: [VertexAiFacetValue]
+}
+
+type VertexAiFacetValue {
+    value: String!
+    count: Int!
+}
+
 type VertexAiAutocompleteResult {
     suggestions: [String]
 }

src/GraphQL/ASTNodeQuery.php

@@ -2,14 +2,14 @@
 
 namespace Drupal\graphql_vertex_ai\GraphQL;
 
-use GraphQL\Language\AST\FieldNode;
+use GraphQL\Language\AST\Node;
 
 /**
  * Navigate the AST\Node architecture to get useful information.
  */
 class ASTNodeQuery {
 
-  public function __construct(protected FieldNode $node) {}
+  public function __construct(protected Node $node) {}
 
   /**
    * Search the AST node hierarchy for a field.
@@ -21,34 +21,72 @@ class ASTNodeQuery {
    *   Whether we found the path or not.
    */
   public function hasField(string $path): bool {
+    $node = $this->nodeAtPath($path);
+    return !empty($node);
+  }
+
+  /**
+   * Get the value of the named parameter at the node path.
+   *
+   * @param string $path
+   *   The path of the node.
+   * @param string $parameter
+   *   The name of the parameter.
+   *
+   * @return string|null
+   *   The value if found.
+   */
+  public function getParameter(string $path, string $parameter): ?string {
+    $node = $this->nodeAtPath($path);
+    if (!$node) {
+      return NULL;
+    }
+    foreach ($node->arguments as $argument) {
+      if ($argument->name->value === $parameter) {
+        return $argument->value->value;
+      }
+    }
+    return NULL;
+  }
+
+  /**
+   * The AST Node at the given path.
+   *
+   * @param string $path
+   *   The path of the node.
+   *
+   * @return \GraphQL\Language\AST\Node|null
+   *   The node if found.
+   */
+  private function nodeAtPath(string $path): ?Node {
     $fields = explode('.', $path);
-    return $this->doHasField($this->node->selectionSet->selections, $fields);
+    return $this->doNodeAtPath($this->node->selectionSet->selections, $fields);
   }
 
   /**
-   * Navigate the tree recursively looking for a field path.
+   * Navigate the tree recursively looking for a node path.
    *
-   * @param \GraphQL\Language\AST\FieldNode[] $nodeList
+   * @param \GraphQL\Language\AST\Node[] $nodeList
    *   The next node in the tree.
    * @param array $fields
    *   The exploded path to search for.
    *
-   * @return bool
-   *   Whether we found the field in the tree or not.
+   * @return \GraphQL\Language\AST\Node|null
+   *   The node if we found it.
    */
-  private function doHasField(iterable $nodeList, array $fields): bool {
+  private function doNodeAtPath(iterable $nodeList, array $fields): ?Node {
     $field = array_shift($fields);
     /** @var \GraphQL\Language\AST\FieldNode $node */
     foreach ($nodeList as $node) {
       if ($node->name->value === $field) {
         if (empty($fields)) {
-          return TRUE;
+          return $node;
         }
         $children = $node->selectionSet?->selections ?? [];
-        return $this->doHasField($children, $fields);
+        return $this->doNodeAtPath($children, $fields);
       }
     }
-    return FALSE;
+    return NULL;
   }
 
 }

src/Plugin/GraphQL/DataProducer/VertexAiSearchQuery.php

@@ -13,6 +13,8 @@ use Google\Cloud\DiscoveryEngine\V1\SearchRequest;
 use Google\Cloud\DiscoveryEngine\V1\SearchRequest\ContentSearchSpec;
 use Google\Cloud\DiscoveryEngine\V1\SearchRequest\ContentSearchSpec\SnippetSpec;
 use Google\Cloud\DiscoveryEngine\V1\SearchRequest\ContentSearchSpec\SummarySpec;
+use Google\Cloud\DiscoveryEngine\V1\SearchRequest\FacetSpec;
+use Google\Cloud\DiscoveryEngine\V1\SearchRequest\FacetSpec\FacetKey;
 use Google\Cloud\DiscoveryEngine\V1\SearchResponse;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
@@ -134,6 +136,21 @@ class VertexAiSearchQuery extends DataProducerPluginBase implements ContainerFac
     $request->setSafeSearch(TRUE);
     $request->setPageSize($pageSize);
     $request->setOffset($offset);
+    if (!empty($filter)) {
+      $request->setFilter($filter);
+    }
+
+    // Turn on facets.
+    $facetParameter = $nodeQuery->getParameter('facets', 'keys');
+    if (!empty($facetParameter)) {
+      $facets = array_map('trim', explode(',', $facetParameter));
+      foreach ($facets as $facet) {
+        $key = (new FacetKey())->setKey($facet);
+        $facetSpec = (new FacetSpec())->setFacetKey($key);
+        $facetSpec->setLimit(50);
+        $request->setFacetSpecs([$facetSpec]);
+      }
+    }
 
     $contentSearchSpec = new ContentSearchSpec();
     // Turn on summary.
@@ -148,10 +165,6 @@ class VertexAiSearchQuery extends DataProducerPluginBase implements ContainerFac
     }
     $request->setContentSearchSpec($contentSearchSpec);
 
-    if (!empty($filter)) {
-      $request->setFilter($filter);
-    }
-
     $response = $searchServiceClient->search($request);
     $results = $response->getPage()->getResponseObject();
 

src/Plugin/GraphQL/SchemaExtension/SearchQuerySchemaExtension.php

@@ -9,6 +9,8 @@ use Drupal\graphql\GraphQL\ResolverRegistryInterface;
 use Drupal\graphql\Plugin\GraphQL\SchemaExtension\SdlSchemaExtensionPluginBase;
 use Drupal\graphql_vertex_ai\GraphQL\Resolver\ASTContext;
 use Google\Cloud\DiscoveryEngine\V1\SearchResponse;
+use Google\Cloud\DiscoveryEngine\V1\SearchResponse\Facet;
+use Google\Cloud\DiscoveryEngine\V1\SearchResponse\Facet\FacetValue;
 use Google\Cloud\DiscoveryEngine\V1\SearchResponse\SearchResult;
 
 /**
@@ -56,6 +58,9 @@ class SearchQuerySchemaExtension extends SdlSchemaExtensionPluginBase {
     $registry->addFieldResolver('VertexAiSearchResponse', 'results',
       $builder->callback(fn (SearchResponse $response) => $response->getResults())
     );
+    $registry->addFieldResolver('VertexAiSearchResponse', 'facets',
+      $builder->callback(fn (SearchResponse $response) => $response->getFacets())
+    );
 
     // Individual search results.
     $registry->addFieldResolver('VertexAiSearchResult', 'id',
@@ -88,7 +93,21 @@ class SearchQuerySchemaExtension extends SdlSchemaExtensionPluginBase {
       )
     );
 
-    // Autocomplete results.
+    // Facets.
+    $registry->addFieldResolver('VertexAiFacet', 'key',
+      $builder->callback(fn (Facet $facet) => $facet->getKey())
+    );
+    $registry->addFieldResolver('VertexAiFacet', 'values',
+      $builder->callback(fn (Facet $facet) => $facet->getValues())
+    );
+    $registry->addFieldResolver('VertexAiFacetValue', 'value',
+      $builder->callback(fn (FacetValue $facetValue) => $facetValue->getValue())
+    );
+    $registry->addFieldResolver('VertexAiFacetValue', 'count',
+      $builder->callback(fn (FacetValue $facetValue) => $facetValue->getCount())
+    );
+
+    // Autocomplete.
     $registry->addFieldResolver('VertexAiAutocompleteResult', 'suggestions',
       $builder->produce('vertex_ai_autocomplete_response_suggestions')
         ->map('response', $builder->fromParent())