-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #155 from miurla/retrieve
Add Tool to Retrieve Page Content from Specific URLs
- Loading branch information
Showing
10 changed files
with
153 additions
and
29 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } | ||
| }) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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> |