Add tavily to search agent (#1166)

Author: wadabee

docs/en/DEPLOY_OPTION.md

@@ -405,12 +405,15 @@ const envs: Record<string, Partial<StackInput>> = {
 
 Creates an Agent that connects to APIs to reference the latest information for responses. You can customize the Agent to add other actions and create multiple Agents to switch between.
 
-The default search agent uses [Brave Search API's Data for AI](https://brave.com/search/api/) due to its large free tier, request limit considerations, and cost factors, but you can customize it to use other APIs. Getting an API key requires credit card registration even for the free plan.
+The default search agents available are [Data for AI in Brave Search API] (https://brave.com/search/api/) or [Tavily's Tavily Search API] (https://docs.tavily.com/documentation/api-reference/endpoint/search). It is also possible to customise the API so that it can be used with other APIs. Please note that Brave Search API requires credit card setup, even for free plans.
 
 > [!NOTE]
-> When you enable the Agent Chat use case, it only sends data to external APIs in the Agent Chat use case. (By default, Brave Search API) Other use cases can continue to be used entirely within AWS. Please check your internal policies and API terms of service before enabling.
+> When you enable the Agent Chat use case, it only sends data to external APIs in the Agent Chat use case. (By default, Brave Search API or Tavily Search) Other use cases can continue to be used entirely within AWS. Please check your internal policies and API terms of service before enabling.
 
-Set `agentEnabled` and `searchAgentEnabled` to `true` (default is `false`), and specify the search engine API key in `searchApiKey`.
+Set `agentEnabled` and `searchAgentEnabled` to `true` (default is `false`), and then set the required fields.
+
+- `searchEngine` : Specify the search engine to use. You can use `Brave` or `Tavily`.
+- `searchApiKey` : Specify the API key of the search engine.
 
 **Edit [parameter.ts](/packages/cdk/parameter.ts)**
 
@@ -420,6 +423,7 @@ const envs: Record<string, Partial<StackInput>> = {
   dev: {
     agentEnabled: true,
     searchAgentEnabled: true,
+    searchEngine: 'Brave' or 'Tavily',
     searchApiKey: '<Search Engine API Key>',
   },
 };
@@ -433,6 +437,7 @@ const envs: Record<string, Partial<StackInput>> = {
   "context": {
     "agentEnabled": true,
     "searchAgentEnabled": true,
+    "searchEngine": "Brave" or "Tavily",
     "searchApiKey": "<Search Engine API Key>"
   }
 }

docs/ja/DEPLOY_OPTION.md

@@ -420,12 +420,15 @@ const envs: Record<string, Partial<StackInput>> = {
 
 API と連携し最新情報を参照して回答する Agent を作成します。Agent のカスタマイズを行い他のアクションを追加できるほか、複数の Agent を作成し切り替えることが可能です。
 
-デフォルトで使用できる検索エージェントでは、無料利用枠の大きさ・リクエスト数の制限・コストの観点から [Brave Search API の Data for AI](https://brave.com/search/api/) を使用していますが、他の API にカスタマイズすることも可能です。API キーの取得はフリープランでもクレジットカードの登録が必要になります。
+デフォルトで使用できる検索エージェントでは、 [Brave Search API の Data for AI](https://brave.com/search/api/) か [Tavily の Tavily Search API](https://docs.tavily.com/documentation/api-reference/endpoint/search) を利用します。他の API を利用するようにカスタマイズすることも可能です。Brave Search API は、無料プランでもクレジットカードの設定が必要なのでご注意ください。
 
 > [!NOTE]
-> Agent チャットユースケースを有効化すると Agent チャットユースケースでのみ外部 API にデータを送信します。（デフォルトでは Brave Search API）他のユースケースは引き続き AWS 内のみに閉じて利用することが可能です。社内ポリシー、API の利用規約などを確認してから有効化してください。
+> Agent チャットユースケースを有効化すると Agent チャットユースケースでのみ外部 API にデータを送信します。（デフォルトでは Brave Search API か Tavily Search API）他のユースケースは引き続き AWS 内のみに閉じて利用することが可能です。社内ポリシー、API の利用規約などを確認してから有効化してください。
 
-`agentEnabled` と `searchAgentEnabled` に `true` を指定し(デフォルトは `false`)、`searchApiKey` に検索エンジンの API キーを指定します。
+`agentEnabled` と `searchAgentEnabled` に `true` (デフォルトは `false`) を指定した上で、必要な項目を設定してください。
+
+- `searchEngine` : 利用する検索エンジンを指定してください。`Brave` か `Tavily` が利用できます。
+- `searchApiKey` : 検索エンジンの API キーを指定します。
 
 **[parameter.ts](/packages/cdk/parameter.ts) を編集**
 
@@ -435,6 +438,7 @@ const envs: Record<string, Partial<StackInput>> = {
   dev: {
     agentEnabled: true,
     searchAgentEnabled: true,
+    searchEngine: 'Brave' or 'Tavily',
     searchApiKey: '<検索エンジンの API キー>',
   },
 };
@@ -448,6 +452,7 @@ const envs: Record<string, Partial<StackInput>> = {
   "context": {
     "agentEnabled": true,
     "searchAgentEnabled": true,
+    "searchEngine": "Brave" or "Tavily",
     "searchApiKey": "<検索エンジンの API キー>"
   }
 }

packages/cdk/cdk.json

@@ -56,6 +56,7 @@
     "endpointNames": [],
     "agentEnabled": false,
     "searchAgentEnabled": false,
+    "searchEngine": "Brave",
     "searchApiKey": "",
     "agents": [],
     "inlineAgents": false,

packages/cdk/lambda/agent.ts

@@ -2,7 +2,69 @@ import {
   AgentInput,
   AgentOutput,
   BraveSearchResult,
+  TavilySearchResult,
 } from 'generative-ai-use-cases';
+import { StackInput } from '../lib/stack-input';
+
+type SearchResult = {
+  title: string;
+  content: string;
+  url: string;
+  extraSnippets?: string[];
+};
+
+const searchUsingBrave = async (keyword: string): Promise<SearchResult[]> => {
+  // https://api-dashboard.search.brave.com/app/documentation/web-search/get-started
+  const searchUrl = `https://api.search.brave.com/res/v1/web/search?q=${keyword}&count=3&text_decorations=0`;
+  const searchApiKey = process.env.SEARCH_API_KEY || '';
+  const response = await fetch(searchUrl, {
+    headers: {
+      'X-Subscription-Token': searchApiKey,
+    },
+  });
+  const data = await response.json();
+  console.log(JSON.stringify(data));
+
+  return data.web.results.map(
+    (result: BraveSearchResult): SearchResult => ({
+      title: result.title,
+      content: result.description,
+      url: result.url,
+      extraSnippets: result.extra_snippets,
+    })
+  );
+};
+
+const searchUsingTavily = async (keyword: string): Promise<SearchResult[]> => {
+  const searchUrl = 'https://api.tavily.com/search';
+  const searchApiKey = process.env.SEARCH_API_KEY || '';
+
+  // https://docs.tavily.com/documentation/api-reference/endpoint/search
+  const response = await fetch(searchUrl, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${searchApiKey}`,
+    },
+    body: JSON.stringify({
+      query: keyword,
+      search_depth: 'basic',
+      include_answer: false,
+      include_images: false,
+      include_raw_content: true,
+      max_results: 3,
+    }),
+  });
+
+  const data = await response.json();
+  console.log(JSON.stringify(data));
+
+  return data.results.map((result: TavilySearchResult) => ({
+    title: result.title,
+    content: result.raw_content ?? result.content,
+    url: result.url,
+  }));
+};
 
 export const handler = async (event: AgentInput): Promise<AgentOutput> => {
   try {
@@ -15,23 +77,13 @@ export const handler = async (event: AgentInput): Promise<AgentOutput> => {
       }
     }
 
-    // Search
-    const searchUrl = `https://api.search.brave.com/res/v1/web/search?q=${keyword}&count=3&text_decorations=0`;
-    const searchApiKey = process.env.SEARCH_API_KEY || '';
-    const response = await fetch(searchUrl, {
-      headers: {
-        'X-Subscription-Token': searchApiKey,
-      },
-    });
-    const data = await response.json();
-    console.log(JSON.stringify(data));
+    const searchEngine = process.env
+      .SEARCH_ENGINE as StackInput['searchEngine'];
 
-    const results = data.web.results.map((result: BraveSearchResult) => ({
-      title: result.title,
-      description: result.description,
-      extra_snippets: result.extra_snippets,
-      url: result.url,
-    }));
+    const results =
+      searchEngine === 'Brave'
+        ? await searchUsingBrave(keyword)
+        : await searchUsingTavily(keyword);
 
     // Create Response Object
     const response_body = {

packages/cdk/lib/agent-stack.ts

@@ -14,11 +14,12 @@ export class AgentStack extends Stack {
   constructor(scope: Construct, id: string, props: AgentStackProps) {
     super(scope, id, props);
 
-    const { searchAgentEnabled, searchApiKey } = props.params;
+    const { searchAgentEnabled, searchApiKey, searchEngine } = props.params;
 
     const agent = new Agent(this, 'Agent', {
       searchAgentEnabled,
       searchApiKey,
+      searchEngine,
     });
 
     this.agents = agent.agents;

packages/cdk/lib/construct/agent.ts

@@ -17,11 +17,13 @@ import { BucketDeployment, Source } from 'aws-cdk-lib/aws-s3-deployment';
 import { CfnAgent, CfnAgentAlias } from 'aws-cdk-lib/aws-bedrock';
 import { Agent as AgentType } from 'generative-ai-use-cases';
 import { LAMBDA_RUNTIME_NODEJS } from '../../consts';
+import { StackInput } from '../stack-input';
 
 interface AgentProps {
   // Context Params
   readonly searchAgentEnabled: boolean;
   readonly searchApiKey?: string | null;
+  readonly searchEngine?: StackInput['searchEngine'];
 }
 
 export class Agent extends Construct {
@@ -32,7 +34,7 @@ export class Agent extends Construct {
 
     const suffix = Lazy.string({ produce: () => Names.uniqueId(this) });
 
-    const { searchAgentEnabled, searchApiKey } = props;
+    const { searchAgentEnabled, searchApiKey, searchEngine } = props;
 
     // Bucket to store schema and data for agents for bedrock
     const s3Bucket = new Bucket(this, 'Bucket', {
@@ -76,7 +78,7 @@ export class Agent extends Construct {
     });
 
     // Search Agent
-    if (searchAgentEnabled && searchApiKey) {
+    if (searchAgentEnabled && searchApiKey && searchEngine) {
       const bedrockAgentLambda = new NodejsFunction(
         this,
         'BedrockAgentLambda',
@@ -86,6 +88,7 @@ export class Agent extends Construct {
           timeout: Duration.seconds(300),
           environment: {
             SEARCH_API_KEY: searchApiKey ?? '',
+            SEARCH_ENGINE: searchEngine,
           },
         }
       );

packages/cdk/lib/stack-input.ts

@@ -1,7 +1,7 @@
 import { z } from 'zod';
 
-// Common Validator
-export const stackInputSchema = z.object({
+// Base schema without refine
+const baseStackInputSchema = z.object({
   account: z.string().default(process.env.CDK_DEFAULT_ACCOUNT ?? ''),
   region: z.string().default(process.env.CDK_DEFAULT_REGION ?? 'us-east-1'),
   env: z.string().default(''),
@@ -125,6 +125,7 @@ export const stackInputSchema = z.object({
   agentEnabled: z.boolean().default(false),
   searchAgentEnabled: z.boolean().default(false),
   searchApiKey: z.string().nullish(),
+  searchEngine: z.enum(['Brave', 'Tavily']).default('Brave'),
   agents: z
     .array(
       z.object({
@@ -164,8 +165,23 @@ export const stackInputSchema = z.object({
   dashboard: z.boolean().default(false),
 });
 
+// Common Validator with refine
+export const stackInputSchema = baseStackInputSchema.refine(
+  (data) => {
+    // If searchApiKey is provided, searchEngine must also be provided
+    if (data.searchApiKey && !data.searchEngine) {
+      return false;
+    }
+    return true;
+  },
+  {
+    message: 'searchEngine is required when searchApiKey is provided',
+    path: ['searchEngine'],
+  }
+);
+
 // schema after conversion
-export const processedStackInputSchema = stackInputSchema.extend({
+export const processedStackInputSchema = baseStackInputSchema.extend({
   modelIds: z.array(
     z.object({
       modelId: z.string(),

packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap

@@ -1109,6 +1109,7 @@ exports[`GenerativeAiUseCases matches the snapshot 3`] = `
         "Environment": {
           "Variables": {
             "SEARCH_API_KEY": "XXXXXX",
+            "SEARCH_ENGINE": "Brave",
           },
         },
         "Handler": "index.handler",

packages/cdk/test/generative-ai-use-cases.test.ts

@@ -47,6 +47,7 @@ describe('GenerativeAiUseCases', () => {
       endpointNames: [],
       agentEnabled: true,
       searchAgentEnabled: true,
+      searchEngine: 'Brave',
       searchApiKey: 'XXXXXX',
       agents: [],
       flows: [],

packages/types/src/agent.d.ts

@@ -36,3 +36,11 @@ export type BraveSearchResult = {
   description: string;
   extra_snippets?: string[];
 };
+
+export type TavilySearchResult = {
+  title: string;
+  url: string;
+  content: string;
+  score: number;
+  raw_content?: string;
+};