From c830a97a6f5313c41b9d9412a68e0abe6e10c2c3 Mon Sep 17 00:00:00 2001
From: Yoshiki Miura <miurap400@gmail.com>
Date: Mon, 17 Jun 2024 10:54:21 +0900
Subject: [PATCH] fix: Add generated answer section component

---
 components/answer-section-generated.tsx | 18 ++++++++++++++++++
 lib/agents/researcher.tsx               | 10 +++++++++-
 lib/agents/writer.tsx                   |  8 +++++++-
 3 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 components/answer-section-generated.tsx

diff --git a/components/answer-section-generated.tsx b/components/answer-section-generated.tsx
new file mode 100644
index 000000000..c3b03b27d
--- /dev/null
+++ b/components/answer-section-generated.tsx
@@ -0,0 +1,18 @@
+'use client'
+
+import { Section } from './section'
+import { BotMessage } from './message'
+
+export type AnswerSectionProps = {
+  result: string
+}
+
+export function AnswerSectionGenerated({ result }: AnswerSectionProps) {
+  return (
+    <div>
+      <Section title="Answer">
+        <BotMessage content={result} />
+      </Section>
+    </div>
+  )
+}
diff --git a/lib/agents/researcher.tsx b/lib/agents/researcher.tsx
index 2d38cd3ec..c1dbc0746 100644
--- a/lib/agents/researcher.tsx
+++ b/lib/agents/researcher.tsx
@@ -3,6 +3,7 @@ import { CoreMessage, ToolCallPart, ToolResultPart, streamText } from 'ai'
 import { getTools } from './tools'
 import { getModel, transformToolMessages } from '../utils'
 import { AnswerSection } from '@/components/answer-section'
+import { AnswerSectionGenerated } from '@/components/answer-section-generated'
 
 export async function researcher(
   uiStream: ReturnType<typeof createStreamableUI>,
@@ -43,7 +44,14 @@ export async function researcher(
     tools: getTools({
       uiStream,
       fullResponse
-    })
+    }),
+    onFinish: event => {
+      // If the response is generated, update the generated answer section
+      // There is a bug where a new instance of the answer section is displayed once when the next section is added
+      if (event.text.length > 0) {
+        uiStream.update(<AnswerSectionGenerated result={event.text} />)
+      }
+    }
   }).catch(err => {
     hasError = true
     fullResponse = 'Error: ' + err.message
diff --git a/lib/agents/writer.tsx b/lib/agents/writer.tsx
index 5e2291d3e..ca9c074f7 100644
--- a/lib/agents/writer.tsx
+++ b/lib/agents/writer.tsx
@@ -2,6 +2,7 @@ import { createStreamableUI, createStreamableValue } from 'ai/rsc'
 import { CoreMessage, streamText } from 'ai'
 import { createOpenAI } from '@ai-sdk/openai'
 import { AnswerSection } from '@/components/answer-section'
+import { AnswerSectionGenerated } from '@/components/answer-section-generated'
 
 export async function writer(
   uiStream: ReturnType<typeof createStreamableUI>,
@@ -28,7 +29,12 @@ export async function writer(
     Link format: [link text](url)
     Image format: ![alt text](url)
     `,
-    messages
+    messages,
+    onFinish: event => {
+      // If the response is generated, update the generated answer section
+      // There is a bug where a new instance of the answer section is displayed once when the next section is added
+      uiStream.update(<AnswerSectionGenerated result={event.text} />)
+    }
   })
     .then(async result => {
       for await (const text of result.textStream) {