fix wrong update

Author: leehuwuj

templates/components/multiagent/python/app/api/routers/vercel_response.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 import logging
-from typing import AsyncGenerator, Awaitable, Generator, List
+from typing import AsyncGenerator, Awaitable, List
 
 from aiostream import stream
 from app.api.routers.models import ChatData, Message
@@ -67,11 +67,6 @@ async def _chat_response_generator():
                 async for token in result:
                     final_response += str(token.delta)
                     yield self.convert_text(token.delta)
-            elif isinstance(result, Generator):
-                for chunk in result:
-                    chunk_str = str(chunk)
-                    final_response += chunk_str
-                    yield self.convert_text(chunk_str)
             else:
                 if hasattr(result, "response"):
                     content = result.response.message.content

templates/types/streaming/fastapi/app/api/routers/vercel_response.py

@@ -1,101 +1,124 @@
-import asyncio
 import json
 import logging
-from typing import AsyncGenerator, Awaitable, List
+from typing import Awaitable, List
 
 from aiostream import stream
-from fastapi import Request
+from fastapi import BackgroundTasks, Request
 from fastapi.responses import StreamingResponse
+from llama_index.core.chat_engine.types import StreamingAgentChatResponse
+from llama_index.core.schema import NodeWithScore
 
-from app.api.routers.models import ChatData, Message
+from app.api.routers.events import EventCallbackHandler
+from app.api.routers.models import ChatData, Message, SourceNodes
 from app.api.services.suggestion import NextQuestionSuggestion
 
 logger = logging.getLogger("uvicorn")
 
 
 class VercelStreamResponse(StreamingResponse):
     """
-    Base class to convert the response from the chat engine to the streaming format expected by Vercel
+    Class to convert the response from the chat engine to the streaming format expected by Vercel
     """
 
     TEXT_PREFIX = "0:"
     DATA_PREFIX = "8:"
 
-    def __init__(self, request: Request, chat_data: ChatData, *args, **kwargs):
-        self.request = request
-        self.chat_data = chat_data
-        content = self.content_generator(*args, **kwargs)
+    def __init__(
+        self,
+        request: Request,
+        event_handler: EventCallbackHandler,
+        response: Awaitable[StreamingAgentChatResponse],
+        chat_data: ChatData,
+        background_tasks: BackgroundTasks,
+    ):
+        content = VercelStreamResponse.content_generator(
+            request, event_handler, response, chat_data, background_tasks
+        )
         super().__init__(content=content)
 
-    async def content_generator(self, event_handler, events):
-        stream = self._create_stream(
-            self.request, self.chat_data, event_handler, events
+    @classmethod
+    async def content_generator(
+        cls,
+        request: Request,
+        event_handler: EventCallbackHandler,
+        response: Awaitable[StreamingAgentChatResponse],
+        chat_data: ChatData,
+        background_tasks: BackgroundTasks,
+    ):
+        chat_response_generator = cls._chat_response_generator(
+            response, background_tasks, event_handler, chat_data
         )
+        event_generator = cls._event_generator(event_handler)
+
+        # Merge the chat response generator and the event generator
+        combine = stream.merge(chat_response_generator, event_generator)
         is_stream_started = False
-        try:
-            async with stream.stream() as streamer:
-                async for output in streamer:
-                    if not is_stream_started:
-                        is_stream_started = True
-                        # Stream a blank message to start the stream
-                        yield self.convert_text("")
-
-                    yield output
-        except asyncio.CancelledError:
-            logger.info("Stopping workflow")
-            await event_handler.cancel_run()
-        except Exception as e:
-            logger.error(
-                f"Unexpected error in content_generator: {str(e)}", exc_info=True
-            )
-        finally:
-            logger.info("The stream has been stopped!")
+        async with combine.stream() as streamer:
+            async for output in streamer:
+                if not is_stream_started:
+                    is_stream_started = True
+                    # Stream a blank message to start displaying the response in the UI
+                    yield cls.convert_text("")
 
-    def _create_stream(
-        self,
-        request: Request,
+                yield output
+
+                if await request.is_disconnected():
+                    break
+
+    @classmethod
+    async def _event_generator(cls, event_handler: EventCallbackHandler):
+        """
+        Yield the events from the event handler
+        """
+        async for event in event_handler.async_event_gen():
+            event_response = event.to_response()
+            if event_response is not None:
+                yield cls.convert_data(event_response)
+
+    @classmethod
+    async def _chat_response_generator(
+        cls,
+        response: Awaitable[StreamingAgentChatResponse],
+        background_tasks: BackgroundTasks,
+        event_handler: EventCallbackHandler,
         chat_data: ChatData,
-        event_handler: Awaitable,
-        events: AsyncGenerator,
-        verbose: bool = True,
     ):
-        # Yield the text response
-        async def _chat_response_generator():
-            result = await event_handler
-            final_response = ""
-
-            if isinstance(result, AsyncGenerator):
-                async for token in result:
-                    final_response += str(token.delta)
-                    yield self.convert_text(token.delta)
-            else:
-                if hasattr(result, "response"):
-                    content = result.response.message.content
-                    if content:
-                        for token in content:
-                            final_response += str(token)
-                            yield self.convert_text(token)
-
-            # Generate next questions if next question prompt is configured
-            question_data = await self._generate_next_questions(
-                chat_data.messages, final_response
-            )
-            if question_data:
-                yield self.convert_data(question_data)
+        """
+        Yield the text response and source nodes from the chat engine
+        """
+        # Wait for the response from the chat engine
+        result = await response
+
+        # Once we got a source node, start a background task to download the files (if needed)
+        cls._process_response_nodes(result.source_nodes, background_tasks)
+
+        # Yield the source nodes
+        yield cls.convert_data(
+            {
+                "type": "sources",
+                "data": {
+                    "nodes": [
+                        SourceNodes.from_source_node(node).model_dump()
+                        for node in result.source_nodes
+                    ]
+                },
+            }
+        )
 
-            # TODO: stream sources
+        final_response = ""
+        async for token in result.async_response_gen():
+            final_response += token
+            yield cls.convert_text(token)
 
-        # Yield the events from the event handler
-        async def _event_generator():
-            async for event in events:
-                event_response = event.to_response()
-                if verbose:
-                    logger.debug(event_response)
-                if event_response is not None:
-                    yield self.convert_data(event_response)
+        # Generate next questions if next question prompt is configured
+        question_data = await cls._generate_next_questions(
+            chat_data.messages, final_response
+        )
+        if question_data:
+            yield cls.convert_data(question_data)
 
-        combine = stream.merge(_chat_response_generator(), _event_generator())
-        return combine
+        # the text_generator is the leading stream, once it's finished, also finish the event stream
+        event_handler.is_done = True
 
     @classmethod
     def convert_text(cls, token: str):
@@ -108,6 +131,24 @@ def convert_data(cls, data: dict):
         data_str = json.dumps(data)
         return f"{cls.DATA_PREFIX}[{data_str}]\n"
 
+    @staticmethod
+    def _process_response_nodes(
+        source_nodes: List[NodeWithScore],
+        background_tasks: BackgroundTasks,
+    ):
+        try:
+            # Start background tasks to download documents from LlamaCloud if needed
+            from app.engine.service import LLamaCloudFileService  # type: ignore
+
+            LLamaCloudFileService.download_files_from_nodes(
+                source_nodes, background_tasks
+            )
+        except ImportError:
+            logger.debug(
+                "LlamaCloud is not configured. Skipping post processing of nodes"
+            )
+            pass
+
     @staticmethod
     async def _generate_next_questions(chat_history: List[Message], response: str):
         questions = await NextQuestionSuggestion.suggest_next_questions(