feat: Support RAG in public preview

PiperOrigin-RevId: 635991192

Author: happy-qiao

src/functions/generate_content.ts

@@ -43,6 +43,8 @@ import {
   formatContentRequest,
   validateGenerateContentRequest,
   validateGenerationConfig,
+  hasVertexRagStore,
+  getApiVersion,
 } from './pre_fetch_processing';
 
 export async function generateContent(
@@ -75,12 +77,13 @@ export async function generateContent(
   };
   const response: Response | undefined = await postRequest({
     region: location,
-    resourcePath: resourcePath,
+    resourcePath,
     resourceMethod: constants.GENERATE_CONTENT_METHOD,
     token: await token,
     data: generateContentRequest,
-    apiEndpoint: apiEndpoint,
-    requestOptions: requestOptions,
+    apiEndpoint,
+    requestOptions,
+    apiVersion: getApiVersion(request),
   }).catch(e => {
     throw new GoogleGenerativeAIError('exception posting request to model', e);
   });
@@ -126,12 +129,13 @@ export async function generateContentStream(
   };
   const response = await postRequest({
     region: location,
-    resourcePath: resourcePath,
+    resourcePath,
     resourceMethod: constants.STREAMING_GENERATE_CONTENT_METHOD,
     token: await token,
     data: generateContentRequest,
-    apiEndpoint: apiEndpoint,
-    requestOptions: requestOptions,
+    apiEndpoint,
+    requestOptions,
+    apiVersion: getApiVersion(request),
   }).catch(e => {
     throw new GoogleGenerativeAIError('exception posting request', e);
   });

src/functions/pre_fetch_processing.ts

@@ -18,8 +18,11 @@
 import {
   GenerateContentRequest,
   GenerationConfig,
+  RetrievalTool,
   SafetySetting,
+  Tool,
 } from '../types/content';
+import {ClientError} from '../types/errors';
 import * as constants from '../util/constants';
 
 export function formatContentRequest(
@@ -55,6 +58,12 @@ export function validateGenerateContentRequest(
       }
     }
   }
+
+  if (hasVertexAISearch(request) && hasVertexRagStore(request)) {
+    throw new ClientError(
+      'Found both vertexAiSearch and vertexRagStore field are set in tool. Either set vertexAiSearch or vertexRagStore.'
+    );
+  }
 }
 
 export function validateGenerationConfig(
@@ -67,3 +76,31 @@ export function validateGenerationConfig(
   }
   return generationConfig;
 }
+
+export function getApiVersion(
+  request: GenerateContentRequest
+): 'v1' | 'v1beta1' {
+  return hasVertexRagStore(request) ? 'v1beta1' : 'v1';
+}
+
+export function hasVertexRagStore(request: GenerateContentRequest): boolean {
+  for (const tool of request?.tools ?? []) {
+    const retrieval = (tool as RetrievalTool).retrieval;
+    if (!retrieval) continue;
+    if (retrieval.vertexRagStore) {
+      return true;
+    }
+  }
+  return false;
+}
+
+export function hasVertexAISearch(request: GenerateContentRequest): boolean {
+  for (const tool of request?.tools ?? []) {
+    const retrieval = (tool as RetrievalTool).retrieval;
+    if (!retrieval) continue;
+    if (retrieval.vertexAiSearch) {
+      return true;
+    }
+  }
+  return false;
+}

src/functions/test/functions_test.ts

@@ -47,6 +47,13 @@ const TEST_USER_CHAT_MESSAGE = [
   {role: constants.USER_ROLE, parts: [{text: TEST_CHAT_MESSAGE_TEXT}]},
 ];
 
+const CONTENTS = [
+  {
+    role: 'user',
+    parts: [{text: 'What is the weater like in Boston?'}],
+  },
+];
+
 const TEST_USER_CHAT_MESSAGE_WITH_GCS_FILE = [
   {
     role: constants.USER_ROLE,
@@ -208,6 +215,12 @@ const TEST_TOOLS_WITH_FUNCTION_DECLARATION: Tool[] = [
   },
 ];
 
+const TEST_TOOLS_WITH_RAG: Tool[] = [
+  {
+    retrieval: {vertexRagStore: {ragResources: [{ragCorpus: 'ragCorpus'}]}},
+  },
+];
+
 const fetchResponseObj = {
   status: 200,
   statusText: 'OK',
@@ -520,9 +533,7 @@ describe('generateContent', () => {
 
   it('returns a FunctionCall when passed a FunctionDeclaration', async () => {
     const req: GenerateContentRequest = {
-      contents: [
-        {role: 'user', parts: [{text: 'What is the weater like in Boston?'}]},
-      ],
+      contents: CONTENTS,
       tools: TEST_TOOLS_WITH_FUNCTION_DECLARATION,
     };
     const expectedResult: GenerateContentResult = {
@@ -559,12 +570,7 @@ describe('generateContent', () => {
 
   it('returns a empty FunctionCall list when response contains invalid data', async () => {
     const req: GenerateContentRequest = {
-      contents: [
-        {
-          role: 'user',
-          parts: [{text: 'What is the weater like in Boston?'}],
-        },
-      ],
+      contents: CONTENTS,
       tools: TEST_TOOLS_WITH_FUNCTION_DECLARATION,
     };
     const expectedResult: GenerateContentResult = {
@@ -594,12 +600,7 @@ describe('generateContent', () => {
 
   it('returns empty candidates when response is empty', async () => {
     const req: GenerateContentRequest = {
-      contents: [
-        {
-          role: 'user',
-          parts: [{text: 'What is the weater like in Boston?'}],
-        },
-      ],
+      contents: CONTENTS,
       tools: TEST_TOOLS_WITH_FUNCTION_DECLARATION,
     };
     fetchSpy.and.resolveTo(new Response(JSON.stringify({}), fetchResponseObj));
@@ -613,6 +614,39 @@ describe('generateContent', () => {
     );
     expect(actualResult.response.candidates).not.toBeDefined();
   });
+
+  it('should use v1 apiVersion', async () => {
+    const request: GenerateContentRequest = {
+      contents: CONTENTS,
+    };
+    fetchSpy.and.resolveTo(buildFetchResponse(TEST_MODEL_RESPONSE));
+    await generateContent(
+      TEST_LOCATION,
+      TEST_RESOURCE_PATH,
+      TEST_TOKEN_PROMISE,
+      request,
+      TEST_API_ENDPOINT
+    );
+    const vertexEndpoint = fetchSpy.calls.allArgs()[0][0];
+    expect(vertexEndpoint).toContain('/v1/');
+  });
+
+  it('should use v1beta1 apiVersion when set RAG in tools', async () => {
+    const request: GenerateContentRequest = {
+      contents: CONTENTS,
+      tools: TEST_TOOLS_WITH_RAG,
+    };
+    fetchSpy.and.resolveTo(buildFetchResponse(TEST_MODEL_RESPONSE));
+    await generateContent(
+      TEST_LOCATION,
+      TEST_RESOURCE_PATH,
+      TEST_TOKEN_PROMISE,
+      request,
+      TEST_API_ENDPOINT
+    );
+    const vertexEndpoint = fetchSpy.calls.allArgs()[0][0];
+    expect(vertexEndpoint).toContain('/v1beta1/');
+  });
 });
 
 describe('generateContentStream', () => {

src/functions/test/pre_fetch_processing_test.ts

@@ -0,0 +1,55 @@
+import {Tool} from '../../types/content';
+import {
+  getApiVersion,
+  validateGenerateContentRequest,
+} from '../pre_fetch_processing';
+
+const TOOL1 = {retrieval: {vertexAiSearch: {datastore: 'datastore'}}} as Tool;
+const TOOL2 = {
+  retrieval: {vertexRagStore: {ragResources: [{ragCorpus: 'ragCorpus'}]}},
+} as Tool;
+const TOOL3 = {
+  retrieval: {
+    vertexAiSearch: {datastore: 'datastore'},
+    vertexRagStore: {ragResources: [{ragCorpus: 'ragCorpus'}]},
+  },
+} as Tool;
+
+const VALID_TOOL_ERROR_MESSAGE =
+  '[VertexAI.ClientError]: Found both vertexAiSearch and vertexRagStore field are set in tool. Either set vertexAiSearch or vertexRagStore.';
+
+describe('validateTools', () => {
+  it('should pass validation when set tool correctly', () => {
+    expect(() =>
+      validateGenerateContentRequest({tools: [TOOL1], contents: []})
+    ).not.toThrow();
+    expect(() =>
+      validateGenerateContentRequest({tools: [TOOL2], contents: []})
+    ).not.toThrow();
+  });
+
+  it('should throw error when set VertexAiSearch and VertexRagStore in two tools in request', () => {
+    expect(() =>
+      validateGenerateContentRequest({tools: [TOOL1, TOOL2], contents: []})
+    ).toThrowError(VALID_TOOL_ERROR_MESSAGE);
+  });
+
+  it('should throw error when set VertexAiSearch and VertexRagStore in a single tool in request', () => {
+    expect(() =>
+      validateGenerateContentRequest({tools: [TOOL3], contents: []})
+    ).toThrowError(VALID_TOOL_ERROR_MESSAGE);
+  });
+});
+
+describe('getApiVersion', () => {
+  it('should return v1', () => {
+    expect(getApiVersion({contents: [], tools: [TOOL1]})).toEqual('v1');
+  });
+
+  it('should return v1beta1', () => {
+    expect(getApiVersion({contents: [], tools: [TOOL2]})).toEqual('v1beta1');
+    expect(getApiVersion({contents: [], tools: [TOOL1, TOOL2]})).toEqual(
+      'v1beta1'
+    );
+  });
+});

src/types/content.ts

@@ -51,7 +51,8 @@ export declare interface VertexInit {
 export declare interface GenerateContentRequest extends BaseModelParams {
   /** Array of {@link Content}.*/
   contents: Content[];
-  /** Optional. The user provided system instructions for the model.
+  /**
+   * Optional. The user provided system instructions for the model.
    * Note: only text should be used in parts of {@link Content}
    */
   systemInstruction?: string | Content;
@@ -110,7 +111,8 @@ export declare interface GetGenerativeModelParams extends ModelParams {
   tools?: Tool[];
   /** Optional. The request options to use for generation. */
   requestOptions?: RequestOptions;
-  /** Optional. The user provided system instructions for the model.
+  /**
+   * Optional. The user provided system instructions for the model.
    * Note: only text should be used in parts of {@link Content}
    */
   systemInstruction?: string | Content;
@@ -138,7 +140,8 @@ export declare interface BaseModelParams {
   generationConfig?: GenerationConfig;
   /** Optional. Array of {@link Tool}. */
   tools?: Tool[];
-  /** Optional. The user provided system instructions for the model.
+  /**
+   * Optional. The user provided system instructions for the model.
    * Note: only text should be used in parts of {@link Content}
    */
   systemInstruction?: string | Content;
@@ -563,12 +566,20 @@ export declare interface CitationMetadata {
  *   date).
  */
 export declare interface GoogleDate {
-  /** Year of the date. Must be from 1 to 9999, or 0 to specify a date without a year. */
+  /**
+   * Year of the date. Must be from 1 to 9999, or 0 to specify a date without a
+   * year.
+   */
   year?: number;
-  /** Month of the date. Must be from 1 to 12, or 0 to specify a year without a monthi and day. */
+  /**
+   * Month of the date. Must be from 1 to 12, or 0 to specify a year without a
+   * monthi and day.
+   */
   month?: number;
-  /** Day of the date. Must be from 1 to 31 and valid for the year and month.
-   * or 0 to specify a year by itself or a year and month where the day isn't significant
+  /**
+   * Day of the date. Must be from 1 to 31 and valid for the year and month.
+   * or 0 to specify a year by itself or a year and month where the day isn't
+   * significant
    */
   day?: number;
 }
@@ -763,6 +774,40 @@ export declare interface RetrievalTool {
   retrieval?: Retrieval;
 }
 
+export declare interface VertexRagStore {
+  /**
+   * Optional. List of corpora for retrieval. Currently only support one corpus
+   * or multiple files from one corpus. In the future we may open up multiple
+   * corpora support.
+   */
+  ragResources?: RagResource[];
+
+  /** Optional. Number of top k results to return from the selected corpora. */
+  similarityTopK?: number;
+
+  /** Optional. If set this field, results with vector distance smaller than this threshold will be returned. */
+  vectorDistanceThreshold?: number;
+}
+
+/**
+ * Config of Vertex RagStore grounding checking.
+ */
+export declare interface RagResource {
+  /**
+   * Optional. Vertex RAG Store corpus resource name.
+   *
+   * @example
+   * `projects/{project}/locations/{location}/ragCorpora/{rag_corpus}`
+   */
+  ragCorpus?: string;
+
+  /**
+   * Optional. Set this field to select the files under the ragCorpora for
+   * retrieval.
+   */
+  ragFileIds?: string[];
+}
+
 /**
  * Defines a retrieval tool that model can call to access external knowledge.
  */
@@ -786,6 +831,10 @@ export declare interface Retrieval {
    * VertexAISearch}.
    */
   vertexAiSearch?: VertexAISearch;
+
+  /** Optional. Set to use data source powered by Vertex RAG store. */
+  vertexRagStore?: VertexRagStore;
+
   /**
    * Optional. Disable using the result from this tool in detecting grounding
    * attribution. This does not affect how the result is given to the model for
@@ -896,7 +945,8 @@ export declare interface StartChatParams {
   tools?: Tool[];
   /** Optional. The base Vertex AI endpoint to use for the request. */
   apiEndpoint?: string;
-  /** Optional. The user provided system instructions for the model.
+  /**
+   * Optional. The user provided system instructions for the model.
    * Note: only text should be used in parts of {@link Content}
    */
   systemInstruction?: string | Content;
@@ -916,7 +966,8 @@ export declare interface StartChatSessionRequest extends StartChatParams {
   publisherModelEndpoint: string;
   /** The resource path to use for the request. */
   resourcePath: string;
-  /** Optional. The user provided system instructions for the model.
+  /**
+   * Optional. The user provided system instructions for the model.
    * Note: only text should be used in parts of {@link Content}
    */
   systemInstruction?: string | Content;