Merge pull request #155 from miurla/retrieve

Add Tool to Retrieve Page Content from Specific URLs

.env.local.example

@@ -31,4 +31,8 @@ UPSTASH_REDIS_REST_TOKEN=[YOUR_UPSTASH_REDIS_REST_TOKEN]
 
 # enable the share feature
 # If you enable this feature, separate account management implementation is required.
-# ENABLE_SHARE=true
+# ENABLE_SHARE=true
+
+# use the retrieve tool
+# Exa API Key retrieved here: https://dashboard.exa.ai/api-keys
+# EXA_API_KEY=[YOUR_EXA_API_KEY]

README.md

@@ -20,6 +20,7 @@ Please note that there are differences between this repository and the official
 - [x] Enable specifying the model to use (only writer agent)
 - [x] Implement search history functionality
 - [x] Develop features for sharing results
+- [x] Implement functionality to get answers from specified URL
 - [ ] Add video support for search functionality
 - [ ] Implement RAG support
 - [ ] Introduce tool support for enhanced productivity
@@ -30,7 +31,7 @@ Please note that there are differences between this repository and the official
 - App framework: [Next.js](https://nextjs.org/)
 - Text streaming / Generative UI: [Vercel AI SDK](https://sdk.vercel.ai/docs)
 - Generative Model: [OpenAI](https://openai.com/)
-- Search API: [Tavily AI](https://tavily.com/)
+- Search API: [Tavily AI](https://tavily.com/) / [Exa AI](https://exa.ai/)
 - Serverless Database: [Upstash](https://upstash.com/)
 - Component library: [shadcn/ui](https://ui.shadcn.com/)
 - Headless component primitives: [Radix UI](https://www.radix-ui.com/)

app/actions.tsx

@@ -20,6 +20,7 @@ import { BotMessage } from '@/components/message'
 import { SearchSection } from '@/components/search-section'
 import SearchRelated from '@/components/search-related'
 import { CopilotDisplay } from '@/components/copilot-display'
+import RetrieveSection from '@/components/retrieve-section'
 
 async function submit(formData?: FormData, skip?: boolean) {
   'use server'
@@ -122,7 +123,7 @@ async function submit(formData?: FormData, skip?: boolean) {
     while (
       useSpecificAPI
         ? toolOutputs.length === 0 && answer.length === 0
-        : answer.length === 0
+        : answer.length === 0 && !errorOccurred
     ) {
       // Search the web and generate the answer
       const { fullResponse, hasError, toolResponses } = await researcher(
@@ -391,6 +392,12 @@ export const getUIStateFromAIState = (aiState: Chat) => {
                   component: <SearchSection result={searchResults.value} />,
                   isCollapsed: isCollapsed.value
                 }
+              case 'retrieve':
+                return {
+                  id,
+                  component: <RetrieveSection data={toolOutput} />,
+                  isCollapsed: isCollapsed.value
+                }
             }
           } catch (error) {
             return {

components/retrieve-section.tsx

@@ -0,0 +1,18 @@
+import React from 'react'
+import { Section } from '@/components/section'
+import { SearchResults } from '@/components/search-results'
+import { SearchResults as SearchResultsType } from '@/lib/types'
+
+interface RetrieveSectionProps {
+  data: SearchResultsType
+}
+
+const RetrieveSection: React.FC<RetrieveSectionProps> = ({ data }) => {
+  return (
+    <Section title="Sources">
+      <SearchResults results={data.results} />
+    </Section>
+  )
+}
+
+export default RetrieveSection

components/search-results.tsx

@@ -28,7 +28,9 @@ export function SearchResults({ results }: SearchResultsProps) {
           <Link href={result.url} passHref target="_blank">
             <Card className="flex-1">
               <CardContent className="p-2">
-                <p className="text-xs line-clamp-2">{result.content}</p>
+                <p className="text-xs line-clamp-2">
+                  {result.title || result.content}
+                </p>
                 <div className="mt-2 flex items-center space-x-2">
                   <Avatar className="h-4 w-4">
                     <AvatarImage

lib/agents/researcher.tsx

@@ -31,7 +31,7 @@ export async function researcher(
   )
 
   let isFirstToolResponse = true
-  const currentDate = new Date().toLocaleString();
+  const currentDate = new Date().toLocaleString()
   const result = await nonexperimental_streamText({
     model: openai.chat(process.env.OPENAI_API_MODEL || 'gpt-4o'),
     maxTokens: 2500,
@@ -46,7 +46,6 @@ export async function researcher(
     tools: getTools({
       uiStream,
       fullResponse,
-      hasError,
       isFirstToolResponse
     })
   })
@@ -73,9 +72,12 @@ export async function researcher(
         break
       case 'tool-result':
         // Append the answer section if the specific model is not used
-        if (!useSpecificModel && toolResponses.length === 0) {
+        if (!useSpecificModel && toolResponses.length === 0 && delta.result) {
           uiStream.append(answerSection)
         }
+        if (!delta.result) {
+          hasError = true
+        }
         toolResponses.push(delta)
         break
       case 'error':

lib/agents/tools/index.tsx

@@ -1,23 +1,33 @@
-import { searchTool } from './search'
 import { createStreamableUI } from 'ai/rsc'
+import { retrieveTool } from './retrieve'
+import { searchTool } from './search'
 
-interface GetToolsProps {
+export interface ToolsProps {
   uiStream: ReturnType<typeof createStreamableUI>
   fullResponse: string
-  hasError: boolean
   isFirstToolResponse: boolean
 }
 
 export const getTools = ({
   uiStream,
   fullResponse,
-  hasError,
   isFirstToolResponse
-}: GetToolsProps) => ({
-  search: searchTool({
-    uiStream,
-    fullResponse,
-    hasError,
-    isFirstToolResponse
-  })
-})
+}: ToolsProps) => {
+  const tools: any = {
+    search: searchTool({
+      uiStream,
+      fullResponse,
+      isFirstToolResponse
+    })
+  }
+
+  if (process.env.EXA_API_KEY) {
+    tools.retrieve = retrieveTool({
+      uiStream,
+      fullResponse,
+      isFirstToolResponse
+    })
+  }
+
+  return tools
+}

lib/agents/tools/retrieve.tsx

@@ -0,0 +1,78 @@
+import { retrieveSchema } from '@/lib/schema/retrieve'
+import { ToolsProps } from '.'
+import { Card } from '@/components/ui/card'
+import { SearchSkeleton } from '@/components/search-skeleton'
+import { SearchResults as SearchResultsType } from '@/lib/types'
+import Exa from 'exa-js'
+import RetrieveSection from '@/components/retrieve-section'
+
+export const retrieveTool = ({
+  uiStream,
+  fullResponse,
+  isFirstToolResponse
+}: ToolsProps) => ({
+  description: 'Retrieve content from the web',
+  parameters: retrieveSchema,
+  execute: async ({ urls }: { urls: string[] }) => {
+    let hasError = false
+    const apiKey = process.env.EXA_API_KEY
+    const exa = new Exa(apiKey)
+
+    // If this is the first tool response, remove spinner
+    if (isFirstToolResponse) {
+      isFirstToolResponse = false
+      uiStream.update(null)
+    }
+    // Append the search section
+    uiStream.append(<SearchSkeleton />)
+
+    let results: SearchResultsType | undefined
+    try {
+      const data = await exa.getContents(urls)
+
+      if (data.results.length === 0) {
+        hasError = true
+      } else {
+        results = {
+          results: data.results.map((result: any) => ({
+            title: result.title,
+            content: result.text,
+            url: result.url
+          })),
+          query: '',
+          images: []
+        }
+      }
+    } catch (error) {
+      hasError = true
+      console.error('Retrieve API error:', error)
+
+      fullResponse += `\n${error} "${urls.join(', ')}".`
+
+      uiStream.update(
+        <Card className="p-4 mt-2 text-sm">
+          {`${error} "${urls.join(', ')}".`}
+        </Card>
+      )
+      return results
+    }
+
+    if (hasError || !results) {
+      fullResponse += `\nAn error occurred while retrieving "${urls.join(
+        ', '
+      )}".`
+      uiStream.update(
+        <Card className="p-4 mt-2 text-sm">
+          {`An error occurred while retrieving "${urls.join(
+            ', '
+          )}".This webiste may not be supported.`}
+        </Card>
+      )
+      return results
+    }
+
+    uiStream.update(<RetrieveSection data={results} />)
+
+    return results
+  }
+})

lib/agents/tools/search.tsx

@@ -1,22 +1,15 @@
-import { createStreamableUI, createStreamableValue } from 'ai/rsc'
+import { createStreamableValue } from 'ai/rsc'
 import Exa from 'exa-js'
 import { searchSchema } from '@/lib/schema/search'
 import { Card } from '@/components/ui/card'
 import { SearchSection } from '@/components/search-section'
-
-interface searchToolProps {
-  uiStream: ReturnType<typeof createStreamableUI>
-  fullResponse: string
-  hasError: boolean
-  isFirstToolResponse: boolean
-}
+import { ToolsProps } from '.'
 
 export const searchTool = ({
   uiStream,
   fullResponse,
-  hasError,
   isFirstToolResponse
-}: searchToolProps) => ({
+}: ToolsProps) => ({
   description: 'Search the web for information',
   parameters: searchSchema,
   execute: async ({
@@ -28,6 +21,7 @@ export const searchTool = ({
     max_results: number
     search_depth: 'basic' | 'advanced'
   }) => {
+    let hasError = false
     // If this is the first tool response, remove spinner
     if (isFirstToolResponse) {
       isFirstToolResponse = false

lib/schema/retrieve.tsx

@@ -0,0 +1,8 @@
+import { DeepPartial } from 'ai'
+import { z } from 'zod'
+
+export const retrieveSchema = z.object({
+  urls: z.array(z.string().url()).describe('The urls to retrieve')
+})
+
+export type PartialInquiry = DeepPartial<typeof retrieveSchema>