fix not display

Author: leehuwuj

templates/components/agents/python/blog/app/workflows/models.py

@@ -1,9 +1,10 @@
 from typing import Any, Dict, List, Literal
 
-from app.api.routers.models import SourceNodes
-from llama_index.core.schema import Node, NodeWithScore
+from llama_index.core.schema import NodeWithScore
 from llama_index.core.workflow import Event
 
+from app.api.routers.models import SourceNodes
+
 
 # Workflow events
 class PlanResearchEvent(Event):
@@ -13,7 +14,7 @@ class PlanResearchEvent(Event):
 class ResearchEvent(Event):
     question_id: str
     question: str
-    context_nodes: List[NodeWithScore | Node]
+    context_nodes: List[NodeWithScore]
 
 
 class CollectAnswersEvent(Event):

templates/components/agents/python/blog/app/workflows/writer.py

@@ -3,16 +3,6 @@
 import uuid
 from typing import Any, Dict, List, Optional
 
-from app.engine.index import IndexConfig, get_index
-from app.workflows.agents import plan_research, research, write_report
-from app.workflows.models import (
-    CollectAnswersEvent,
-    DataEvent,
-    PlanResearchEvent,
-    ResearchEvent,
-    SourceNodesEvent,
-    WriteReportEvent,
-)
 from llama_index.core.indices.base import BaseIndex
 from llama_index.core.memory import ChatMemoryBuffer
 from llama_index.core.memory.simple_composable_memory import SimpleComposableMemory
@@ -26,7 +16,18 @@
     step,
 )
 
-logger = logging.getLogger(__name__)
+from app.engine.index import IndexConfig, get_index
+from app.workflows.agents import plan_research, research, write_report
+from app.workflows.models import (
+    CollectAnswersEvent,
+    DataEvent,
+    PlanResearchEvent,
+    ResearchEvent,
+    SourceNodesEvent,
+    WriteReportEvent,
+)
+
+logger = logging.getLogger("uvicorn")
 logger.setLevel(logging.INFO)
 
 
@@ -45,7 +46,7 @@ def create_workflow(
     return WriterWorkflow(
         index=index,
         chat_history=chat_history,
-        **kwargs,
+        timeout=120.0,
     )
 
 
@@ -67,22 +68,20 @@ class WriterWorkflow(Workflow):
     context_nodes: List[Node]
     index: BaseIndex
     user_request: str
-    stream: bool = False
+    stream: bool = True
 
     def __init__(
         self,
         index: BaseIndex,
         chat_history: Optional[List[ChatMessage]] = None,
-        stream: bool = False,
-        timeout: Optional[float] = 120.0,
+        stream: bool = True,
         **kwargs,
     ):
         super().__init__(**kwargs)
         self.index = index
         self.context_nodes = []
         self.stream = stream
         self.chat_history = chat_history
-        self.timeout = timeout
         self.memory = SimpleComposableMemory.from_defaults(
             primary_memory=ChatMemoryBuffer.from_defaults(
                 chat_history=chat_history,
@@ -142,7 +141,7 @@ async def analyze(
         """
         Analyze the retrieved information
         """
-        print("Analyzing the retrieved information")
+        logger.info("Analyzing the retrieved information")
         ctx.write_event_to_stream(
             DataEvent(
                 type="analyze",
@@ -262,7 +261,7 @@ async def collect_answers(
                     content=f"<Question>{result.question}</Question>\n<Answer>{result.answer}</Answer>",
                 )
             )
-        ctx.set("n_questions", 0)
+        await ctx.set("n_questions", 0)
         self.memory.put(
             message=ChatMessage(
                 role=MessageRole.ASSISTANT,
@@ -276,6 +275,7 @@ async def report(self, ctx: Context, ev: WriteReportEvent) -> StopEvent:
         """
         Report the answers
         """
+        logger.info("Writing the report")
         res = await write_report(
             memory=self.memory,
             user_request=self.user_request,

templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx

@@ -25,7 +25,7 @@ export function ChatMessageContent() {
     },
     // add the writer card
     {
-      position: ContentPosition.AFTER_EVENTS,
+      position: ContentPosition.CHAT_EVENTS,
       component: <WriterCard message={message} />,
     },
     {

templates/types/streaming/nextjs/app/components/ui/chat/custom/writer-card.tsx

@@ -8,7 +8,7 @@ import {
   NotebookPen,
   Search,
 } from "lucide-react";
-import { useEffect, useState } from "react";
+import { useMemo } from "react";
 import {
   Collapsible,
   CollapsibleContent,
@@ -45,42 +45,49 @@ type WriterState = {
   };
 };
 
-// Update the state based on the event
-const updateState = (state: WriterState, event: WriterEvent): WriterState => {
+const stateIcon: Record<EventState, React.ReactNode> = {
+  pending: <Clock className="w-4 h-4 text-yellow-500" />,
+  inprogress: <CircleDashed className="w-4 h-4 text-blue-500 animate-spin" />,
+  done: <CheckCircle2 className="w-4 h-4 text-green-500" />,
+  error: <AlertCircle className="w-4 h-4 text-red-500" />,
+};
+
+// Transform the state based on the event without mutations
+const transformState = (
+  state: WriterState,
+  event: WriterEvent,
+): WriterState => {
   switch (event.type) {
     case "answer": {
       const { id, question, answer } = event.data;
       if (!id || !question) return state;
 
-      const questions = state.analyze.questions;
-      const existingQuestion = questions.find((q) => q.id === id);
-
-      const updatedQuestions = existingQuestion
-        ? questions.map((q) =>
-            q.id === id
-              ? {
-                  ...existingQuestion,
-                  state: event.state,
-                  answer: answer || existingQuestion.answer,
-                }
-              : q,
-          )
-        : [
-            ...questions,
+      const updatedQuestions = state.analyze.questions.map((q) => {
+        if (q.id !== id) return q;
+        return {
+          ...q,
+          state: event.state,
+          answer: answer ?? q.answer,
+        };
+      });
+
+      const newQuestion = !state.analyze.questions.some((q) => q.id === id)
+        ? [
             {
               id,
               question,
-              answer: answer || null,
+              answer: answer ?? null,
               state: event.state,
               isOpen: false,
             },
-          ];
+          ]
+        : [];
 
       return {
         ...state,
         analyze: {
           ...state.analyze,
-          questions: updatedQuestions,
+          questions: [...updatedQuestions, ...newQuestion],
         },
       };
     }
@@ -100,36 +107,30 @@ const updateState = (state: WriterState, event: WriterEvent): WriterState => {
   }
 };
 
-export function WriterCard({ message }: { message: Message }) {
-  const [state, setState] = useState<WriterState>({
+// Convert writer events to state
+const writeEventsToState = (events: WriterEvent[] | undefined): WriterState => {
+  if (!events?.length) {
+    return {
+      retrieve: { state: null },
+      analyze: { state: null, questions: [] },
+    };
+  }
+
+  const initialState: WriterState = {
     retrieve: { state: null },
     analyze: { state: null, questions: [] },
-  });
+  };
 
+  return events.reduce(
+    (acc: WriterState, event: WriterEvent) => transformState(acc, event),
+    initialState,
+  );
+};
+
+export function WriterCard({ message }: { message: Message }) {
   const writerEvents = message.annotations as WriterEvent[] | undefined;
 
-  useEffect(() => {
-    if (writerEvents?.length) {
-      writerEvents.forEach((event) => {
-        setState((currentState) => updateState(currentState, event));
-      });
-    }
-  }, [writerEvents]);
-
-  const getStateIcon = (state: EventState | null) => {
-    switch (state) {
-      case "pending":
-        return <Clock className="w-4 h-4 text-yellow-500" />;
-      case "inprogress":
-        return <CircleDashed className="w-4 h-4 text-blue-500 animate-spin" />;
-      case "done":
-        return <CheckCircle2 className="w-4 h-4 text-green-500" />;
-      case "error":
-        return <AlertCircle className="w-4 h-4 text-red-500" />;
-      default:
-        return null;
-    }
-  };
+  const state = useMemo(() => writeEventsToState(writerEvents), [writerEvents]);
 
   if (!writerEvents?.length) {
     return null;
@@ -167,7 +168,7 @@ export function WriterCard({ message }: { message: Message }) {
                   <CollapsibleTrigger className="w-full">
                     <div className="flex items-center gap-2 p-3 hover:bg-gray-50 transition-colors rounded-lg border border-gray-200">
                       <div className="flex-shrink-0">
-                        {getStateIcon(question.state)}
+                        {stateIcon[question.state]}
                       </div>
                       <span className="font-medium text-left flex-1">
                         {question.question}