Merge pull request #183 from nextcloud/feat/streaming

feat: Stream output text and tool calls

Author: marcelklehr

.github/workflows/integration_test.yml

@@ -285,6 +285,42 @@ jobs:
           echo $TASK | jq '.ocs.data.task.output.sources'
           echo $TASK | jq '.ocs.data.task.output.sources' | grep -q 'get_folder_tree'
 
+      - name: Run streaming task
+        if: ${{ matrix.server-versions == 'master' }}
+        env:
+          CREDS: "alice:alice"
+        run: |
+          TASK=$(curl -X POST -u "$CREDS" -H "oCS-APIRequest: true" -H "Content-type: application/json" \
+            "http://localhost:8080/ocs/v2.php/taskprocessing/schedule?format=json" \
+            --data-raw '{"input": {"input": "Count to 500", "confirmation":1, "conversation_token": ""},"type":"core:contextagent:interaction", "appId": "test", "customId": "", "preferStreaming": true}')
+          echo "$TASK"
+          TASK_ID=$(echo "$TASK" | jq '.ocs.data.task.id')
+
+          STREAMING_OBSERVED=0
+          ELAPSED=0
+          TASK_STATUS='"STATUS_SCHEDULED"'
+
+          until [ "$ELAPSED" -ge 80 ] || [ "$TASK_STATUS" = '"STATUS_SUCCESSFUL"' ] || [ "$TASK_STATUS" = '"STATUS_FAILED"' ]; do
+            sleep 2
+            ELAPSED=$((ELAPSED + 2))
+            TASK=$(curl -u "$CREDS" -H "oCS-APIRequest: true" "http://localhost:8080/ocs/v2.php/taskprocessing/task/$TASK_ID?format=json")
+            echo "$TASK"
+            TASK_STATUS=$(echo "$TASK" | jq '.ocs.data.task.status')
+            echo "Status: $TASK_STATUS"
+
+            if [ "$TASK_STATUS" = '"STATUS_RUNNING"' ]; then
+              TASK_OUTPUT=$(echo "$TASK" | jq -r '.ocs.data.task.output.output // empty')
+              if [ -n "$TASK_OUTPUT" ]; then
+                echo "Streaming update observed: $TASK_OUTPUT"
+                STREAMING_OBSERVED=1
+              fi
+            fi
+          done
+
+          curl -u "$CREDS" -H "oCS-APIRequest: true" "http://localhost:8080/ocs/v2.php/taskprocessing/task/$TASK_ID?format=json"
+          [ "$TASK_STATUS" = '"STATUS_SUCCESSFUL"' ]
+          [ "$STREAMING_OBSERVED" -eq 1 ]
+
       - name: Show nextcloud logs
         if: always()
         run: |

ex_app/lib/agent.py

@@ -4,10 +4,13 @@
 import os
 import string
 import random
+from collections.abc import Awaitable, Callable
 from datetime import date
+from time import monotonic
+from typing import Any, cast
 
-from langchain_core.messages import ToolMessage, SystemMessage, AIMessage, HumanMessage
-from langchain_core.runnables import RunnableConfig, RunnableLambda
+from langchain_core.messages import ToolMessage, SystemMessage, AIMessage, HumanMessage, AIMessageChunk
+from langchain_core.runnables import RunnableConfig
 from nc_py_api import AsyncNextcloudApp
 from nc_py_api.ex_app import persistent_storage
 
@@ -94,7 +97,11 @@ def export_conversation(checkpointer):
 	conversation_token = add_signature(serialized_state.decode('utf-8'), key)
 	return conversation_token
 
-async def react(task, nc: AsyncNextcloudApp):
+async def react(
+		task,
+		nc: AsyncNextcloudApp,
+		stream_output: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
+):
 	safe_tools, dangerous_tools = await get_tools(nc)
 
 	model.bind_nextcloud(nc)
@@ -183,14 +190,62 @@ async def call_model(
 	else:
 		new_input = {"messages": [("user", task['input']['input'])]}
 
-	async for event in graph.astream(new_input, thread, stream_mode="values"):
+	snapshot_messages = state_snapshot.values.get('messages', [])
+	last_message: AIMessage = AIMessage("")
+	if len(snapshot_messages) > 0:
+		last_message = cast(AIMessage, snapshot_messages[-1])
+	previous_message_count = len(snapshot_messages)
+	source_list: list[str] = []
+	known_sources: set[str] = set()
+	streamed_output = ''
+	last_stream_update = 0.0
+	last_reported_stream_state: dict[str, Any] | None = None
+	prefer_streaming = bool(task.get('preferStreaming'))
+	stream_mode = ["messages", "values"] if prefer_streaming and stream_output is not None else "values"
+
+	async def report_stream_state(force: bool = False):
+		nonlocal last_stream_update
+		nonlocal last_reported_stream_state
+		if stream_output is None:
+			return
+		stream_state = {'output': streamed_output, 'sources': json.dumps(source_list.copy())}
+		if last_reported_stream_state == stream_state:
+			return
+		now = monotonic()
+		if not force and last_reported_stream_state is not None and (now - last_stream_update) < 0.5:
+			return
+		await stream_output(stream_state)
+		last_stream_update = now
+		last_reported_stream_state = stream_state
+
+	async for event in graph.astream(new_input, thread, stream_mode=stream_mode):
+		if isinstance(event, tuple):
+			mode, payload = event
+		else:
+			mode, payload = "values", event
+
+		if mode == 'messages':
+			message_chunk, metadata = payload
+			if metadata.get('langgraph_node') != 'agent' or not isinstance(message_chunk, AIMessageChunk):
+				continue
+			chunk_content = message_chunk.content
+			if isinstance(chunk_content, str) and chunk_content != '':
+				streamed_output += chunk_content
+				await report_stream_state()
+			continue
+
+		event = payload
 		last_message = event['messages'][-1]
-		for message in event['messages']:
-			if isinstance(message, HumanMessage):
-				source_list = []
+		for message in event['messages'][previous_message_count:]:
 			if isinstance(message, AIMessage) and message.tool_calls:
 					for tool_call in message.tool_calls:
-						source_list.append(tool_call['name'])
+						tool_name = tool_call['name']
+						if tool_name not in known_sources:
+							known_sources.add(tool_name)
+							source_list.append(tool_name)
+							await report_stream_state(force=True)
+
+	await report_stream_state(force=True)
 
 	state_snapshot = graph.get_state(thread)
 	actions = ''

ex_app/lib/main.py

@@ -176,7 +176,26 @@ async def handle_task(task, nc: AsyncNextcloudApp):
         nextcloud = AsyncNextcloudApp()
         if task['userId']:
             await nextcloud.set_user(task['userId'])
-        output = await react(task, nextcloud)
+
+        stream_updates_enabled = task.get('preferStreaming', None) is True
+        stream_update_failed = False
+
+        async def stream_output(intermediate_output):
+            nonlocal stream_update_failed
+            if not stream_updates_enabled or stream_update_failed:
+                return
+            try:
+                await nc.ocs(
+                    "POST",
+                    f"/ocs/v2.php/taskprocessing/tasks_provider/{task['id']}/stream-result",
+                    json={"output": intermediate_output},
+                )
+            except (NextcloudException, RequestException) as stream_err:
+                stream_update_failed = True
+                tb_str = ''.join(traceback.format_exception(stream_err))
+                await log(nc, LogLvl.WARNING, "Error streaming intermediate task result: " + tb_str)
+
+        output = await react(task, nextcloud, stream_output=stream_output if stream_updates_enabled else None)
     except Exception as e:  # noqa
         try:
             tb_str = ''.join(traceback.format_exception(e))