diff --git a/README.md b/README.md
index e646e5b6d55aca5cbb216fb0d65ca83f84400213..efa1e8bc716ff491b990ea87533590064eb26f76 100644
--- a/README.md
+++ b/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
diff --git a/examples/autocomplete_query.graphql b/examples/autocomplete_query.graphql
index 374f353a1f7315f3819ba876d23dffc8d294394e..c56cfb9670839f37a23c8ddef2d6ac94d9a38e07 100644
--- a/examples/autocomplete_query.graphql
+++ b/examples/autocomplete_query.graphql
@@ -1,5 +1,5 @@
{
- vertexAiAutocompleteQuery(query:"party") {
+ vertexAiAutocompleteQuery(query: "party") {
suggestions
}
}
diff --git a/examples/search_query.graphql b/examples/search_query.graphql
index 8fb47928079a6dadfb819a314cc44ec72ffb36d8..45b5cb4c2c3061d74c6525bdf018dd7a5aa305da 100644
--- a/examples/search_query.graphql
+++ b/examples/search_query.graphql
@@ -1,5 +1,5 @@
{
- 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
+ }
+ }
}
}
diff --git a/graphql/vertex_ai_search_query.base.graphqls b/graphql/vertex_ai_search_query.base.graphqls
index 3e3bf8f7a2b5ec97b2d75d31560fe65e9621390b..5108d79fecae7511432bd8cb0c77351995d7b4e3 100644
--- a/graphql/vertex_ai_search_query.base.graphqls
+++ b/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]
}
diff --git a/src/GraphQL/ASTNodeQuery.php b/src/GraphQL/ASTNodeQuery.php
index e8d4546efae746a346e7a9dfdceedb0d998885a5..020717dac1a1c601c1c8cb9e6f8d908f022e7e93 100644
--- a/src/GraphQL/ASTNodeQuery.php
+++ b/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;
}
}
diff --git a/src/Plugin/GraphQL/DataProducer/VertexAiSearchQuery.php b/src/Plugin/GraphQL/DataProducer/VertexAiSearchQuery.php
index dc140eaf3dc4d8b4b8af001fe128281053cd5487..3c46ebe0656aa33cb740f5b91278cfa1f36879eb 100644
--- a/src/Plugin/GraphQL/DataProducer/VertexAiSearchQuery.php
+++ b/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(20);
+ $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();
diff --git a/src/Plugin/GraphQL/SchemaExtension/SearchQuerySchemaExtension.php b/src/Plugin/GraphQL/SchemaExtension/SearchQuerySchemaExtension.php
index d8ffdf89d66f41b88066ddf9aab2f2c0365d17fe..331981c2687a64aae7e5c5ee65113a7f5d7140ec 100644
--- a/src/Plugin/GraphQL/SchemaExtension/SearchQuerySchemaExtension.php
+++ b/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())