Callbacks not working as expected

[ChristianEvc]

Hi everyone,

I'm building a chatbot with a python backend and a react native front end, and a websockets pipe to connect the two.

Everything has been working like a charm, but I'm just adding a feature to provide a source from a RAG pipeline, but I cannot get a callback to work for the source specific chain in Langgraph.

At a high level, I'm using langgraph. If the decision is made to do RAG, the retrieved documents are stored in a state variable, and then two other nodes are triggered:

"Gernerate" - This one simply summarises the documents into a response
"Source" - This one reads the retrieved docs, but summarises the metadata and the Chapter / Subchapter info into a JSON object.

The "Generate" one works perfectly and streams each token to the front end with no issues.
The "Source" one, I don't want to stream token by token as this is a 3 line json object, so its quick to produce. However, there doesn't seemt to be any event I can subscribe to create an appropriate callback.

Below are the node definitions, and the callbacks applied to the LLMs for each:

GENERATE (works fine):

`
def generate(state, streaming_llm):
"""
Generate answer
Args:
state (dict): The current graph state

Returns:
    state (dict): New key added to state, generation, that contains LLM generation
"""
print("---GENERATE---")
question = state["single_question"]
documents = state["documents"]
gen_chain = gen_prompt | streaming_llm

# RAG generation
generation = gen_chain.invoke({"context": documents, "question": question})
return {"documents": documents, "question": question, "generation": generation}`

The "streaming_llm" in this case is set up with the following callback handler (using OPEN AIs ChatOpenAI() ):
`class StreamingLLMCallbackHandler(AsyncCallbackHandler):
"""Callback handler for streaming LLM responses."""

def __init__(self, websocket):
    self.websocket = websocket

async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
    resp = ChatResponse(sender="bot", message=token, type="stream")
    await self.websocket.send_json(resp.model_dump())`

"SOURCE" (Doesn't work)
`def sourceHandling(state, source_llm):
"""
Based on the answer generated, determine the relevant source from the retrieved documents.

Args:
    state (dict): The current graph state

Returns:
    state (dict): New key added to state, generation, that contains LLM generation
"""
print("--SOURCE--")
documents = state["documents"]
answer = state["generation"]
source_llm_structured = source_llm.with_structured_output(Source)
source_chain = source_prompt | source_llm_structured
source = source_chain.invoke({"context": documents, "answer": answer})
return {"sourceData": source}`

The "source_llm" in this case is set up with the following callback handler (using OPEN AIs ChatOpenAI() ):
`class SourceCallbackHandler(AsyncCallbackHandler):
"""Async callback handler that can be used to handle callbacks from langchain."""
def init(self, websocket):
self.websocket = websocket

async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
    """Run when chain ends running."""
    print("RAW SOURCE LLM response:", response)
    resp = SourceResponse(sender="bot", message=response, type="source")
    await self.websocket.send_json(resp.model_dump())

async def on_tool_end(self, output: Any, **kwargs: Any) -> None:
    print("---TOOL END---")
    print(output)`

Whilst the "on_llm_end" event is triggering, its not allowing me to do anything useful with the LLMResult response. And no other event seems to even be triggered (for example, the "on_tool_end" event).

Whilst this may seem like a convoluted way to stream the answers back, its been the lowest latency way I've found, and also the only way that allows me to directly stream to websockets.

However, I'm sure there is a better way, so any help or insights would be hugely appreciated!!