Refactor retrieveTool to use jina.ai

miurla · May 18, 2024 · 3db99cb · 3db99cb
1 parent 6d40d74
commit 3db99cb
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 28 deletions.
diff --git a/.env.local.example b/.env.local.example
@@ -31,8 +31,4 @@ 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
-
-# use the retrieve tool
-# Exa API Key retrieved here: https://dashboard.exa.ai/api-keys
-# EXA_API_KEY=[YOUR_EXA_API_KEY]
+# ENABLE_SHARE=true
diff --git a/lib/agents/tools/retrieve.tsx b/lib/agents/tools/retrieve.tsx
@@ -3,7 +3,6 @@ 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 = ({
@@ -13,10 +12,8 @@ export const retrieveTool = ({
 }: ToolsProps) => ({
   description: 'Retrieve content from the web',
   parameters: retrieveSchema,
-  execute: async ({ urls }: { urls: string[] }) => {
+  execute: async ({ url }: { url: 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) {
@@ -28,17 +25,24 @@ export const retrieveTool = ({
 
     let results: SearchResultsType | undefined
     try {
-      const data = await exa.getContents(urls)
-
-      if (data.results.length === 0) {
+      const response = await fetch(`https://r.jina.ai/${url}`, {
+        method: 'GET',
+        headers: {
+          Accept: 'application/json'
+        }
+      })
+      const json = await response.json()
+      if (!json.data || json.data.length === 0) {
         hasError = true
       } else {
         results = {
-          results: data.results.map((result: any) => ({
-            title: result.title,
-            content: result.text,
-            url: result.url
-          })),
+          results: [
+            {
+              title: json.data.title,
+              content: json.data.content,
+              url: json.data.url
+            }
+          ],
           query: '',
           images: []
         }
@@ -47,25 +51,19 @@ export const retrieveTool = ({
       hasError = true
       console.error('Retrieve API error:', error)
 
-      fullResponse += `\n${error} "${urls.join(', ')}".`
+      fullResponse += `\n${error} "${url}".`
 
       uiStream.update(
-        <Card className="p-4 mt-2 text-sm">
-          {`${error} "${urls.join(', ')}".`}
-        </Card>
+        <Card className="p-4 mt-2 text-sm">{`${error} "${url}".`}</Card>
       )
       return results
     }
 
     if (hasError || !results) {
-      fullResponse += `\nAn error occurred while retrieving "${urls.join(
-        ', '
-      )}".`
+      fullResponse += `\nAn error occurred while retrieving "${url}".`
       uiStream.update(
         <Card className="p-4 mt-2 text-sm">
-          {`An error occurred while retrieving "${urls.join(
-            ', '
-          )}".This webiste may not be supported.`}
+          {`An error occurred while retrieving "${url}".This webiste may not be supported.`}
         </Card>
       )
       return results

diff --git a/lib/schema/retrieve.tsx b/lib/schema/retrieve.tsx
@@ -2,7 +2,7 @@ import { DeepPartial } from 'ai'
 import { z } from 'zod'
 
 export const retrieveSchema = z.object({
-  urls: z.array(z.string().url()).describe('The urls to retrieve')
+  url: z.string().url().describe('The url to retrieve')
 })
 
 export type PartialInquiry = DeepPartial<typeof retrieveSchema>