Merge pull request #77 from miurla/refactor-researcher

Fix Researcher
miurla · Apr 24, 2024 · 90332c1 · 90332c1
2 parents fcd93c2 + 0db3e86
commit 90332c1
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 35 deletions.
diff --git a/app/action.tsx b/app/action.tsx
@@ -42,8 +42,6 @@ async function submit(formData?: FormData, skip?: boolean) {
   }
 
   async function processEvents() {
-    uiStream.update(<Spinner />)
-
     let action: any = { object: { next: 'proceed' } }
     // If the user skips the task, we proceed to the search
     if (!skip) action = (await taskManager(messages)) ?? action
@@ -70,6 +68,7 @@ async function submit(formData?: FormData, skip?: boolean) {
     let toolOutputs = []
     let errorOccurred = false
     const streamText = createStreamableValue<string>()
+    uiStream.update(<Spinner />)
 
     // If useSpecificAPI is enabled, only function calls will be made
     // If not using a tool, this model generates the answer

diff --git a/components/search-section.tsx b/components/search-section.tsx
@@ -0,0 +1,44 @@
+'use client'
+
+import { SearchResults } from './search-results'
+import { SearchSkeleton } from './search-skeleton'
+import { SearchResultsImageSection } from './search-results-image'
+import { Section } from './section'
+import { ToolBadge } from './tool-badge'
+import type { SearchResults as TypeSearchResults } from '@/lib/types'
+import { StreamableValue, useStreamableValue } from 'ai/rsc'
+
+export type SearchSectionProps = {
+  result?: StreamableValue<string>
+}
+
+export function SearchSection({ result }: SearchSectionProps) {
+  const [data, error, pending] = useStreamableValue(result)
+  const results: TypeSearchResults = data ? JSON.parse(data) : undefined
+  return (
+    <div>
+      {!pending && data ? (
+        <>
+          <Section size="sm" className="pt-2 pb-0">
+            <ToolBadge tool="search">{`${results.query}`}</ToolBadge>
+          </Section>
+          {results.images && results.images.length > 0 && (
+            <Section title="Images">
+              <SearchResultsImageSection
+                images={results.images}
+                query={results.query}
+              />
+            </Section>
+          )}
+          <Section title="Results">
+            <SearchResults results={results.results} />
+          </Section>
+        </>
+      ) : (
+        <Section className="pt-2 pb-0">
+          <SearchSkeleton />
+        </Section>
+      )}
+    </div>
+  )
+}
diff --git a/lib/agents/researcher.tsx b/lib/agents/researcher.tsx
@@ -8,13 +8,11 @@ import {
 import { searchSchema } from '@/lib/schema/search'
 import { Section } from '@/components/section'
 import { OpenAI } from '@ai-sdk/openai'
-import { ToolBadge } from '@/components/tool-badge'
-import { SearchSkeleton } from '@/components/search-skeleton'
-import { SearchResults } from '@/components/search-results'
 import { BotMessage } from '@/components/message'
 import Exa from 'exa-js'
-import { SearchResultsImageSection } from '@/components/search-results-image'
 import { Card } from '@/components/ui/card'
+import { SearchResults } from '../types'
+import { SearchSection } from '@/components/search-section'
 
 export async function researcher(
   uiStream: ReturnType<typeof createStreamableUI>,
@@ -38,6 +36,7 @@ export async function researcher(
     </Section>
   )
 
+  let isFirstToolResponse = true
   const result = await experimental_streamText({
     model: openai.chat(process.env.OPENAI_API_MODEL || 'gpt-4-turbo'),
     maxTokens: 2500,
@@ -61,17 +60,14 @@ export async function researcher(
           max_results: number
           search_depth: 'basic' | 'advanced'
         }) => {
-          uiStream.update(
-            <Section>
-              <ToolBadge tool="search">{`${query}`}</ToolBadge>
-            </Section>
-          )
-
-          uiStream.append(
-            <Section>
-              <SearchSkeleton />
-            </Section>
-          )
+          // If this is the first tool response, remove spinner
+          if (isFirstToolResponse) {
+            isFirstToolResponse = false
+            uiStream.update(null)
+          }
+          // Append the search section
+          const streamResults = createStreamableValue<string>()
+          uiStream.append(<SearchSection result={streamResults.value} />)
 
           // Tavily API requires a minimum of 5 characters in the query
           const filledQuery =
@@ -97,24 +93,7 @@ export async function researcher(
             return searchResult
           }
 
-          uiStream.update(
-            <Section title="Images">
-              <SearchResultsImageSection
-                images={searchResult.images}
-                query={searchResult.query}
-              />
-            </Section>
-          )
-          uiStream.append(
-            <Section title="Sources">
-              <SearchResults results={searchResult.results} />
-            </Section>
-          )
-
-          // Append the answer section if the specific model is not used
-          if (!useSpecificModel) {
-            uiStream.append(answerSection)
-          }
+          streamResults.done(JSON.stringify(searchResult))
 
           return searchResult
         }
@@ -142,6 +121,10 @@ export async function researcher(
         toolCalls.push(delta)
         break
       case 'tool-result':
+        // Append the answer section if the specific model is not used
+        if (!useSpecificModel && toolResponses.length === 0) {
+          uiStream.append(answerSection)
+        }
         toolResponses.push(delta)
         break
       case 'error':

diff --git a/lib/types/index.ts b/lib/types/index.ts
@@ -0,0 +1,11 @@
+export type SearchResults = {
+  images: string[]
+  results: SearchResultItem[]
+  query: string
+}
+
+export type SearchResultItem = {
+  title: string
+  url: string
+  content: string
+}