Simple Vertexai RAG implementation

DictionaryAgentService

This is the JAVA class that uses the Data Store that was set up to analyse the prompt.

Connection to AI Application

To do the search, the SearchRequest has to be coupled with a Agent which in turn is connected to the datastore. To create the agent one uses the Google Console under the project that will be used (in this application it is blurock-database). Here there is a ‘create app’ button. 

The first decision is the type of app to be created. In our case, under Search and Assistant the Custom Search (general) app was chosen.

The next step is to choose the configuration:

1. Generative Responses: this type of app is important for the search that is intended, namely using a RAG search and some interpretation of the terms that are found.
2. App Name: Here the ‘real’ app name will be this name with a number attached (see below field). This full name (with number attached) is what is used to refer to this app from the JAVA code.
3. Company Name: 
4. Location: In this case, this has to be compatible with the other services (Firestore, Data Store). ‘global’ is a safe bet (unless there are legal restrictions.

The next step is to connect the app with the Data Store. The list of currently available Data Store databases are listed.

The last step is to determine the pricing. And a this point one has an agent that can be referred to by the search.

Connection to App in the JAVA code

The  location of the Data Store to be used in the Agent is given as a combination of project, location, and agent Id (the long name):

String agentId = “analyzeterm_1772469548697”;

String servingConfigS = String.format(

“projects/%s/locations/global/collections/default_collection/engines/%s/servingConfigs/default_search”,

“blurock-database”,

agentId

);

This is used in the search setup (SearchRequest):

.setServingConfig(servingConfigS)

Data Store Filter

There is only one Data Store dictionary that has all the terms. However, there could be different types of terms for different purposes. This is determined by the metadata within the dictionary term, specifically the datatype.

This filter is used to isolate only those dictionary terms having the proper datatype:

String filterExpression = “metadata.datatype: ANY(\”” + datatype + “\”) “;

This is used within the SearchRequest setup with:

.setFilter(filterExpression)

Instruction String

The instruction string of the agent tells the agent how to use the Data Store values. In this experiment, the goal was actually just to identify the task and to identify the name of the task. The task is found with the RAG procedure and the value is determined through the ‘Generative Response’ type that was enable in the configuration. 

The instructions are given to the agent on how to use the RAG and the prompt

ROLE: You are a Data Extraction Assistant.

TASK: You are provided with a USER QUERY and a list of DICTIONARY TERMS (Search Results).

Your goal is to find the best matching ‘term’ from the results and extract the specific value associated with it from the user’s query.

RULES:

1. KEY: Select the ‘term’ from the search results that matches the user’s intent.

2. VALUE: Extract the value/name (e.g., ‘SimpleBenson’) directly from the USER QUERY.

3. OUTPUT: Return only JSON. If no value is found in the query, use null.

USER QUERY: {query}

FORMAT:

{“term_name”: “extracted_value”}

SearchRequest Definition

This configuration is put together with the SearchRequest: 

SearchRequest request = SearchRequest.newBuilder()

.setServingConfig(servingConfigS)

.setQuery(userPrompt)

.setFilter(filterExpression)

.setContentSearchSpec(SearchRequest.ContentSearchSpec.newBuilder()

.setSummarySpec(SearchRequest.ContentSearchSpec.SummarySpec.newBuilder()

.setSummaryResultCount(1)

.setIgnoreLowRelevantContent(true)

.setModelPromptSpec(SearchRequest.ContentSearchSpec.SummarySpec.ModelPromptSpec.newBuilder()

.setPreamble(jsonInstructionString)

.build())

.build())

.build())

.setPageSize(1) // Limit to top 5 identified terms

.build();

Using the SearchRequest for the response

The response is found in the SearchResponse:

SearchPagedResponse pagedResponse = client.search(request);

SearchResponse response = pagedResponse.getPage().getResponse();

System.out.println(response); // The LLM’s interpretation)

Interpreting the Response

Within the SearchResponse, the result is under the summary:

String rawJsonOutput = response.getSummary().getSummaryText();

A typical summary result is:

summary {

summary_text: ““`json\n{\”InitialReadInLocalStorageSystem\”: \”SimpleBenson\”}\n“`”

summary_with_metadata {

summary: ““`json\n{\”InitialReadInLocalStorageSystem\”: \”SimpleBenson\”}\n“`”

citation_metadata {

}

}

}

From this the JSON value should be isolated:

String cleanJson = rawJsonOutput.replaceAll(““`json|“`”, “”).trim();