Emit custom event

Author: hinthornw

backend/retrieval_graph/researcher_graph/graph.py

@@ -10,6 +10,7 @@
 from langchain_core.runnables import RunnableConfig
 from langgraph.constants import Send
 from langgraph.graph import END, START, StateGraph
+from langgraph.config import get_stream_writer
 from typing_extensions import TypedDict
 
 from backend import retrieval
@@ -67,8 +68,10 @@ async def retrieve_documents(
     Returns:
         dict[str, list[Document]]: A dictionary with a 'documents' key containing the list of retrieved documents.
     """
+    writer = get_stream_writer()
     with retrieval.make_retriever(config) as retriever:
         response = await retriever.ainvoke(state.query, config)
+        writer({"retrieve_documents": {"index": state.index, "documents": response}})
         return {"documents": response}
 
 

frontend/app/contexts/GraphContext.tsx

@@ -826,6 +826,9 @@ export function GraphProvider({ children }: { children: ReactNode }) {
       }
     },
     onUpdateEvent: handleUpdateEvent,
+    onCustomEvent: (event) => {
+      console.log("Custom event", event);
+    },
     onFinish: async (state, meta) => {
       setIsStreaming(false);
       const currentState = streamStateRef.current;
@@ -1092,7 +1095,7 @@ export function GraphProvider({ children }: { children: ReactNode }) {
             metadata: userId ? { user_id: userId } : undefined,
             streamResumable: true,
             onDisconnect: "continue",
-            streamMode: ["messages", "messages-tuple", "updates"],
+            streamMode: ["messages-tuple", "custom"],
           },
         );
       } catch (error) {