Rank search results

Select a documentation version:

Learn how to rank and rerank your search results for applications using Vertex AI ranking models such as semantic-ranker-default-003. You can use the ai.rank() function to score documents based on relevance to a query, and improve vector search results by reranking them for better query ordering.

The Vertex AI ranking API takes a list of documents and ranks those documents based on how relevant the documents are to a given query (a search string). When you use the ai.rank() function, it returns scores for how well a document answers a given query.

To use instructions on this page, you must have an understanding of AlloyDB Omni and be familiar with generative AI concepts.

AlloyDB Omni reserves, and creates, the ai schema.

Before you begin

Before you rank search results, do the following:

  1. Configure user access to Vertex AI models.
  2. Verify that the latest version of google_ml_integration is installed.

    1. To check the installed version, run the following command:

      SELECT extversion FROM pg_extension WHERE extname = 'google_ml_integration';
extversion
------------
1.5.2
(1 row)

    2. If the extension isn't installed or if the installed version is earlier than 1.5.2, update the extension.

      CREATE EXTENSION IF NOT EXISTS google_ml_integration;
ALTER EXTENSION google_ml_integration UPDATE;

      If you experience issues when you run the preceding commands, or if the extension isn't updated to version 1.5.2 after you run the preceding commands, contact Google Cloud support.

  3. To use the AlloyDB AI query engine functionality, set the google_ml_integration.enable_ai_query_engine flag to on.

    1. Use a text editor to set the following flag in the postgresql.conf configuration file for your installation of AlloyDB Omni: 
      google_ml_integration.enable_ai_query_engine = on
    2. After you save the postgresql.conf file, restart the AlloyDB Omni service for the changes to take effect. 
      sudo systemctl restart alloydbomni18

  4. Enable the Discover Engine API.

  5. Get the required roles to use ranking models.

Enable the Discovery Engine API

<p>To use ranking models, you must enable the Discovery Engine API.<br>
    Replace <code><var>PROJECT_ID</var></code> with your
    Google Cloud project ID and <code><var>PROJECT_NUMBER</var></code>
    with your corresponding project number.</p>

<pre class="prettyprint lang-terminal">
    # Enable Discovery Engine API
    gcloud services enable discoveryengine.googleapis.com --project=<var>PROJECT_ID</var>
    gcloud projects add-iam-policy-binding <var>PROJECT_ID</var> \
    --member="serviceAccount:service-<var>PROJECT_NUMBER</var>@iam.gserviceaccount.com" \
    --role="roles/discoveryengine.viewer"
</pre>

<p>Model registration for ranking isn't required for Vertex AI models.
    You can use the Vertex AI model name as the
    <code>model_id</code>, which is shown in the following example.
</p>

<pre class="prettyprint lang-sh">
    SELECT index, score
    FROM
      ai.rank(
        model_id => 'semantic-ranker-default-003',
        search_string => 'Affordable family-friendly vacation spots in Southeast Asia?',
        documents =>
    ARRAY[
      'Luxury resorts in South Korea',
      'Family vacation packages for Vietnam: Ha Long Bay and Hoi An',
      'Budget-friendly beaches in Thailand perfect for families',
      'A backpacker guide to solo travel in India'])
</pre>

<p>A common use case for the semantic ranker is to rerank the results returned
    by vector search for better query ordering. The following example shows how to
    use the semantic ranking model for this use case. The example retrieves an
    initial result set for the query <code>personal fitness
    equipment</code> using vector search. These results are then re-ranked to
    return the top five results.
</p>

<pre class="prettyprint lang-sh">
    WITH initial_ranking AS (
      SELECT id, description, ROW_NUMBER() OVER () AS ref_number
      FROM product
      ORDER BY
        embedding <=> google_ml.embedding(
          'gemini-embedding-001', 'personal fitness equipment')::vector
      LIMIT 10
    ), reranked_results AS (
      SELECT index, score
      FROM ai.rank(
          model_id => 'semantic-ranker-default-003',
          search_string => 'personal fitness equipment',
          documents => (SELECT ARRAY_AGG(description ORDER BY ref_number) FROM initial_ranking),
          top_n => 5)
    )
    SELECT id, description
    FROM initial_ranking, reranked_results
    WHERE initial_ranking.ref_number = reranked_results.index
    ORDER BY reranked_results.score DESC;
  </pre>

  <p>For a list of available models and use cases, see <a href="/generative-ai-app-builder/docs/ranking#models">Supported models</a>.</p>

Integrate with Vertex AI and install the extension

For more information, see Configure your AlloyDB Omni installation to query cloud-based models.

Required roles

To get the permissions that you need to use ranking models from Discovery Engine, ask your administrator to grant AlloyDB Omni service account the Discovery Engine Viewer (roles/discoveryengine.viewer) Identity and Access Management (IAM) role on your project. For more information about granting roles, see Manage access to projects, folders, and organizations.

You might also be able to get the required permissions through custom roles or other predefined roles.

Rank your search results

The following SQL query shows how to rank your search results :

SELECT
  ai.rank(
    model_id => 'MODEL_ID',
    search_string => 'SEARCH_STRING',
    documents => ARRAY['DOCUMENT_1', 'DOCUMENT_2', 'DOCUMENT_3']);

Replace the following:

Parameter Description
MODEL_ID A unique ID for the model endpoint that you define.
SEARCH_STRING A search string against which the records are ranked.
DOCUMENTS An array of unique text strings that you want to rank.

For a list of the supported Vertex AI ranking models, see Supported models.

Examples

To rank search results using a Vertex AI ranking model, run the following query:

SELECT index, score
FROM
  ai.rank(
    model_id => 'semantic-ranker-default-003',
    search_string => 'AlloyDB is a PostgreSQL compatible AI database that is ready for production.',
    documents =>
      ARRAY[
        'Alloys are made from combination of metals',
        'The best enterprise-ready PostgreSQL database.',
        'You can feel the heat in Alloy apartments.']);

The response is a table that shows each document and the score based on relevance to the search query. The following is the sample response:

 index | score
-------+------------
     2 |  0.33
     1 |  0.28
     3 |  0.16
(3 rows)

Consider an example AlloyDB Omni database with a list of review descriptions that are converted to embeddings. The following sample code snippet shows how to use the ranking model to retrieve the name of the top-ranked products based on their review descriptions' semantic similarity to a query.

WITH initial_ranking AS (
    SELECT product_id, name, review, review_id, ROW_NUMBER() OVER () AS ref_number
    FROM user_reviews
    ORDER BY
      review_desc_embedding <=> google_ml.embedding(
        'gemini-embedding-001', 'good desserts')::vector
    LIMIT 10
  ), reranked_results AS (
    SELECT index, score
    FROM
      ai.rank(
        model_id => 'semantic-ranker-512',
        search_string => 'good desserts',
        documents => (SELECT ARRAY_AGG(review ORDER BY ref_number) FROM initial_ranking),
        top_n => 5)
  )
SELECT product_id, name
FROM initial_ranking, reranked_results
WHERE initial_ranking.ref_number = reranked_results.index
ORDER BY reranked_results.score DESC;

What's next