From 0f9eda5f6883555a10d47d22a7f184aed1ed2dbf Mon Sep 17 00:00:00 2001 From: Tavily PR Agent Date: Thu, 11 Jun 2026 14:25:57 +0000 Subject: [PATCH] feat: add Tavily as parallel search provider option --- .env.example | 1 + app/api/answer/route.ts | 13 ++++- components/EndpointToggle.tsx | 7 ++- lib/env.ts | 1 + lib/tavily.ts | 60 ++++++++++++++++++++++ lib/types.ts | 4 +- package.json | 1 + tests/setup.ts | 1 + tests/tavily.test.ts | 97 +++++++++++++++++++++++++++++++++++ 9 files changed, 180 insertions(+), 5 deletions(-) create mode 100644 lib/tavily.ts create mode 100644 tests/tavily.test.ts diff --git a/.env.example b/.env.example index da8c18b..eb0de32 100644 --- a/.env.example +++ b/.env.example @@ -1,2 +1,3 @@ BRAVE_API_KEY= # Brave Search API subscription token ANTHROPIC_API_KEY= # Anthropic Console API key (sk-ant-...) +TAVILY_API_KEY= # Tavily Search API key (tvly-...) — optional, needed for Tavily endpoint diff --git a/app/api/answer/route.ts b/app/api/answer/route.ts index 31d8f10..b4cb926 100644 --- a/app/api/answer/route.ts +++ b/app/api/answer/route.ts @@ -1,6 +1,7 @@ import { anthropic } from "@ai-sdk/anthropic"; import { streamText } from "ai"; import { BraveError, llmContext, webSearch } from "@/lib/brave"; +import { TavilyError, tavilySearch } from "@/lib/tavily"; import { logger } from "@/lib/logger"; import { buildPrelude } from "@/lib/prelude"; import { buildContext, buildSystemPrompt } from "@/lib/prompt"; @@ -69,13 +70,21 @@ export async function POST(req: Request): Promise { // 3. Retrieve sources via the selected endpoint. let retrieval: RetrievalResult; try { - retrieval = - endpoint === "web" ? await webSearch(question) : await llmContext(question); + if (endpoint === "tavily") { + retrieval = await tavilySearch(question); + } else { + retrieval = + endpoint === "web" ? await webSearch(question) : await llmContext(question); + } } catch (err) { if (err instanceof BraveError) { logger.error("answer.brave_error", { status: err.status }); return json({ error: "Search provider error." }, 502); } + if (err instanceof TavilyError) { + logger.error("answer.tavily_error", { message: err.message }); + return json({ error: "Search provider error." }, 502); + } logger.error("answer.retrieval_error", { message: err instanceof Error ? err.message : "unknown", }); diff --git a/components/EndpointToggle.tsx b/components/EndpointToggle.tsx index 25f6284..05b4b4d 100644 --- a/components/EndpointToggle.tsx +++ b/components/EndpointToggle.tsx @@ -13,6 +13,11 @@ const OPTIONS: { value: Endpoint; label: string; description: string }[] = [ label: "LLM Context", description: "Model-ready grounding in a single call.", }, + { + value: "tavily", + label: "Tavily", + description: "LLM-optimized search via Tavily API.", + }, ]; export function EndpointToggle({ @@ -28,7 +33,7 @@ export function EndpointToggle({
{OPTIONS.map((opt) => { const active = opt.value === value; diff --git a/lib/env.ts b/lib/env.ts index 919ebb2..ae3616f 100644 --- a/lib/env.ts +++ b/lib/env.ts @@ -3,6 +3,7 @@ import { z } from "zod"; const envSchema = z.object({ BRAVE_API_KEY: z.string().min(1, "BRAVE_API_KEY is required"), ANTHROPIC_API_KEY: z.string().min(1, "ANTHROPIC_API_KEY is required"), + TAVILY_API_KEY: z.string().optional(), // Per-IP rate limit (optional; sensible demo defaults). RATE_LIMIT_MAX: z.coerce.number().int().positive().default(30), RATE_LIMIT_WINDOW_MS: z.coerce.number().int().positive().default(3_600_000), diff --git a/lib/tavily.ts b/lib/tavily.ts new file mode 100644 index 0000000..1a878d6 --- /dev/null +++ b/lib/tavily.ts @@ -0,0 +1,60 @@ +import { tavily } from "@tavily/core"; +import { getEnv } from "@/lib/env"; +import { logger } from "@/lib/logger"; +import type { RetrievalResult } from "@/lib/types"; + +const DEFAULT_MAX_RESULTS = 10; + +/** Thrown when Tavily returns an error or the API key is missing. */ +export class TavilyError extends Error { + constructor(message: string) { + super(message); + this.name = "TavilyError"; + } +} + +/** Tavily Search → normalized sources + retrieval time. */ +export async function tavilySearch(query: string): Promise { + const { TAVILY_API_KEY } = getEnv(); + if (!TAVILY_API_KEY) { + throw new TavilyError("TAVILY_API_KEY is not configured"); + } + + const client = tavily({ apiKey: TAVILY_API_KEY }); + + const start = Date.now(); + let response; + try { + response = await client.search(query, { + maxResults: DEFAULT_MAX_RESULTS, + searchDepth: "basic", + }); + } catch (err) { + const ms = Date.now() - start; + logger.error("tavily.request_error", { + ms, + message: err instanceof Error ? err.message : "unknown", + }); + throw new TavilyError( + err instanceof Error ? err.message : "Tavily search failed", + ); + } + const ms = Date.now() - start; + + logger.info("tavily.request", { + ms, + resultCount: response.results?.length ?? 0, + }); + + const results = response.results ?? []; + const sources = results + .filter((r: { url?: string }) => r.url) + .map((r: { title?: string; url: string; content?: string }, i: number) => ({ + index: i + 1, + title: r.title || r.url, + url: r.url, + snippet: r.content ?? "", + })); + + return { sources, retrievalMs: ms }; +} diff --git a/lib/types.ts b/lib/types.ts index 3094781..ead2d3f 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -10,7 +10,7 @@ export const SourceSchema = z.object({ export type Source = z.infer; /** The retrieval endpoints exposed by the app. */ -export type Endpoint = "web" | "context"; +export type Endpoint = "web" | "context" | "tavily"; /** Result of a single Brave retrieval call: normalized sources + elapsed time. */ export type RetrievalResult = { @@ -33,7 +33,7 @@ export type AnswerPrelude = { /** Request body for POST /api/answer. */ export const AnswerRequestSchema = z.object({ question: z.string().trim().min(1, "question is required").max(400, "question is too long"), - endpoint: z.enum(["web", "context"]), + endpoint: z.enum(["web", "context", "tavily"]), }); export type AnswerRequest = z.infer; diff --git a/package.json b/package.json index 3ecd427..c2799be 100644 --- a/package.json +++ b/package.json @@ -21,6 +21,7 @@ "next": "15.5.18", "react": "^19.2.6", "react-dom": "^19.2.6", + "@tavily/core": "^0.6.1", "zod": "^4.4.3" }, "devDependencies": { diff --git a/tests/setup.ts b/tests/setup.ts index 7bdd12b..f0f7961 100644 --- a/tests/setup.ts +++ b/tests/setup.ts @@ -1,2 +1,3 @@ process.env.BRAVE_API_KEY ??= "test-brave-key"; process.env.ANTHROPIC_API_KEY ??= "test-anthropic-key"; +process.env.TAVILY_API_KEY ??= "test-tavily-key"; diff --git a/tests/tavily.test.ts b/tests/tavily.test.ts new file mode 100644 index 0000000..d09b74a --- /dev/null +++ b/tests/tavily.test.ts @@ -0,0 +1,97 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +// Mock @tavily/core before importing the module under test. +const mockSearch = vi.fn(); +vi.mock("@tavily/core", () => ({ + tavily: () => ({ search: mockSearch }), +})); + +// Import after mock is registered. +const { tavilySearch, TavilyError } = await import("@/lib/tavily"); + +describe("tavilySearch", () => { + beforeEach(() => mockSearch.mockReset()); + afterEach(() => vi.restoreAllMocks()); + + it("normalizes Tavily results into a 1-based Source[]", async () => { + mockSearch.mockResolvedValueOnce({ + results: [ + { title: "T1", url: "https://1.com", content: "Snippet one" }, + { title: "T2", url: "https://2.com", content: "Snippet two" }, + ], + }); + + const result = await tavilySearch("hello"); + expect(result.sources).toEqual([ + { index: 1, title: "T1", url: "https://1.com", snippet: "Snippet one" }, + { index: 2, title: "T2", url: "https://2.com", snippet: "Snippet two" }, + ]); + expect(typeof result.retrievalMs).toBe("number"); + expect(result.retrievalMs).toBeGreaterThanOrEqual(0); + }); + + it("returns empty sources when results array is empty", async () => { + mockSearch.mockResolvedValueOnce({ results: [] }); + + const result = await tavilySearch("empty query"); + expect(result.sources).toEqual([]); + }); + + it("filters out results without a url", async () => { + mockSearch.mockResolvedValueOnce({ + results: [ + { title: "No URL", url: "", content: "text" }, + { title: "Has URL", url: "https://x.com", content: "text" }, + ], + }); + + const result = await tavilySearch("q"); + expect(result.sources).toHaveLength(1); + expect(result.sources[0].url).toBe("https://x.com"); + }); + + it("throws TavilyError when the API key is missing", async () => { + const original = process.env.TAVILY_API_KEY; + delete process.env.TAVILY_API_KEY; + + // Clear the env cache so getEnv() re-reads process.env. + vi.resetModules(); + const freshEnvMod = await import("@/lib/env"); + // Access the cached value to force re-parse — we need to clear it. + // The env module caches on first call; we re-mock to bypass. + const { getEnv } = freshEnvMod; + + // Since env is cached, we test at the tavily module level instead. + // Restore the key and test error path by clearing TAVILY_API_KEY from the + // already-parsed env. The simplest approach: the tavily module checks the + // key from getEnv() which returns the cached Env. We simulate missing key + // by temporarily patching getEnv. + const envMod = await import("@/lib/env"); + const originalGetEnv = envMod.getEnv; + vi.spyOn(envMod, "getEnv").mockReturnValue({ + BRAVE_API_KEY: "test", + ANTHROPIC_API_KEY: "test", + RATE_LIMIT_MAX: 30, + RATE_LIMIT_WINDOW_MS: 3_600_000, + TAVILY_API_KEY: undefined, + } as ReturnType); + + // Re-import tavily to pick up the spy (but the mock is module-level, so + // we call the already-imported function which calls getEnv internally). + await expect(tavilySearch("q")).rejects.toMatchObject({ + name: "TavilyError", + message: "TAVILY_API_KEY is not configured", + }); + + process.env.TAVILY_API_KEY = original; + }); + + it("throws TavilyError when the upstream client throws", async () => { + mockSearch.mockRejectedValueOnce(new Error("API rate limit exceeded")); + + await expect(tavilySearch("q")).rejects.toMatchObject({ + name: "TavilyError", + message: "API rate limit exceeded", + }); + }); +});