RunnableWithMessageHistory: add ability to erase session content

[pavel604]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Thanks for adding the RunnableWithMessageHistory class. It really helps to keep the ducks in a row. However, it appears that there's no way to delete the history by session ID, e.g. after user has exited the chat. So it just grows and grows until the agent is restarted. Also, would be nice to fix the token tracking for the agents. It seems to only work on simple LLM run calls.

Motivation

This, hopefully, solves two problems:

• fix session clutter in an agent
• fix broken token callback in agents (LLMResult doesn't return token count or price)
  Thanks!

Proposal (If applicable)

No response

[eyurtsev]

Thanks for the suggestion!

Session clearing does not belong in RunnableWithMessageHistory:

1. the runnable has no way of knowing when the session has ended
2. Clearing history should not require a user to produce a RunnableWithMessageHistory (which requires an underlying runnable to be instantiated)

I would recommend one of the following two options, depending on the goal:

1. For clearing the completely history -- use the underlying persistence layer abstraction.
2. If history trimming is needed, it can be pre-pended as a RunnableLambda to the Runnable that is being wrapped by the RunnableWithMessageHistory.

---

Fix broken token callback in agents (LLMResult doesn't return token count or price)

Could you say a bit more about this issue? This feels like an unrelated issue, but something we'd want to fix.

[pavel604]

Regarding item 1: of course, the server controls the session, and what I'm trying to do is:

My agent:

    _llm = ChatOpenAI(temperature=0, 
                      model='gpt-4-1106-preview', 
                      streaming=True, 
                      verbose=verbose,
                      callback_manager=BaseCallbackManager(
                          handlers=[Fb()]),
                    #   callbacks=[Fb()]
                      )
    print(u'\rLLM launched.')
    prompt = hub.pull("hwchase17/openai-functions-agent")
    
    tools.append(PythonREPLTool())
    
    _agent = create_openai_functions_agent(_llm, tools, prompt)
    _agent_executor = AgentExecutor(agent=_agent, 
                                    tools=tools, 
                                    verbose=verbose,
                                    callbacks=[Fb()]
                                    )
    # return _agent_executor
    
    _message_history = ChatMessageHistory()
    
    agent = RunnableWithMessageHistory(
        _agent_executor,
        lambda session_id: _message_history,
        input_messages_key="input",
        history_messages_key="chat_history",
    )

Then, when the server sees the user leave:

def on_user_left_chat(session_id):
    """Called when user exits the chatroom"""
    db.save_thread(agent.get_session_history(session_id=session_id)) # works
    agent.forget_conversation(session_id) # <--- the missing method

This way I can have several users having concurrent conversations and an easy way to manage the session histories.

Regarding item 2, If you run the agent code above, with the Fb callback being:

class Fb(BaseCallbackHandler):

    _instance = None
    def __new__(cls, *args, **kwargs):
        if not isinstance(cls._instance, cls):
            print('Creating new Firebase interface object')
            cls._instance = super(Fb, cls).__new__(cls)
            # Put any initialization here.
        return cls._instance
    
    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
        """Run when LLM ends running."""
        print(response.llm_output)
        print(response.run)
        print(response.generations)

The llm_output and run are none. My understanding is that either field in LLMResult should contain token count.

Even if I do this:

    with get_openai_callback() as cb:
        result = agent.invoke(
            { "input": query },
            config={"configurable": {"session_id": session_id}},
        )["output"]
    
        print(cb)

I get nothing in cb.

I would appreciate an advice on the latter, as I've tried several variants and the only thing that returns token count is if you call _llm.run()

Thanks!