How to Best Implement LLM Summaries of Streamed Values in LangChain?

[nick-youngblut]

I’m working on a project where I use LangChain to create a react agent. I want to generate concise summaries of each step in the agent’s workflow as it streams. Currently, I’m using a separate LLM chain to handle the summaries, but I’m wondering if there’s a more efficient or better way to implement this.

Here’s a simplified version of my code:

from langgraph.prebuilt import create_react_agent
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Function to create the summary chain
def create_step_summary_chain(model_name: str = "gpt-4o-mini", max_tokens: int = 45):
    template = (
        "Concisely summarize the following step in the workflow:\n\n"
        "{step}\n\n"
        "Summary (max {max_tokens} tokens):"
    )
    prompt = PromptTemplate(
        input_variables=["step"],
        template=template.format(max_tokens=max_tokens)
    )
    llm = OpenAI(model_name=model_name, temperature=0, max_tokens=max_tokens)
    return LLMChain(llm=llm, prompt=prompt)

# Function to invoke the agent and generate summaries
def invoke_agent(agent, input_message: str, step_summary_chain=None):
    for i, step in enumerate(agent.stream(input_message)):
        if step_summary_chain:
            summary = step_summary_chain.run(step=step)
            print(f"Step {i+1} Summary: {summary}")
        else:
            print(f"Step {i+1}: {step}")

# Example usage
if __name__ == "__main__":
    # Create the agent (simplified for this example)
    agent = create_react_agent(...)

    # Create the step summary chain
    step_summary_chain = create_step_summary_chain()

    # Invoke the agent with input and summary chain
    invoke_agent(agent, input_message="Do some task", step_summary_chain=step_summary_chain)

Example output:

Step 1: Summary of the react agent's first message
Step 2: Summary of the react agent's second message
...

One problem with this approach is that langsmith tracks each invocation of step_summary_chain. I've tried to disable langsmith tracking of step_summary_chain invocations, but I haven't found a solution.