Pass custom context to callbacks

[tubedude]

For context, I'm using LangChain to analise large documents all within one LLMChain. The messages I receive are treated differently at each step of the analysis.
The problem I'm facing (besides the fact I'm struggling to understand the JS way of implementing callbacks) is the fact that the custom_context is not available in ChatModel callbacks, as it is only available in the LLMChain callbacks.

So the idea here would be to make custom_context available to all callbacks. What do you think?

My use case would look like this:

defmodule MyApp.DocAnalyzer do
  def handlers() do
    %{
    on_llm_new_delta: &handle_llm_new_delta/3
    }
  end
  
  def handle_llm_new_delta(llmchain, _index, delta) do
    %{custom_context: %{liveview: liveview_pid, step: step}} = llmchain
    send(liveview_pid, {:new_delta, step, delta}
  end
  
  def init(liveview_pid, document) do
    llmchain =
      [llm: ChatOpenAI.new!()]
      |> LLMChain.new!()
      |> LLMChain.update_custom_context(%{liveview: liveview_pid, step: :none})
      |> LLMChain.add_message(Message.new_user!(document)
  end

  def execute_step_one(llmchain, data) do
     # This step will not trigger the call back as it would set stream to false.
     
     {:ok, llmchain, response} =
       llmchain
       |> LLMChain.add_message(Message.new_user!("Give me a JSON of doc features.")
       |> LLMChain.update_custom_context(%{step: :one})
       |> LLMChain.run(%{stream: false})
  end

  def execute_step_two(llmchain, data) do
     {:ok, llmchain, response} =
       llmchain
       |> LLMChain.add_message(Message.new_user!("Fill in the abstract")
       |> LLMChain.update_custom_context(%{step: :two})
       |> LLMChain.run(%{stream: true})
  end
  
    def execute_step_three(llmchain, data) do
     {:ok, llmchain, response} =
       llmchain
       |> LLMChain.add_message(Message.new_user!("Fill in the other part...")
       |> LLMChain.update_custom_context(%{step: :three})
       |> LLMChain.run(%{stream: true})
  end
end

... and those functions would be called from a LiveView at each user interaction...

On a more conceptual discussion, I'm not understanding why the stream related callbacks are under Model callbacks and not under Chain callbacks, as it feels new_delta is something generated within the Chain, and not the Model in itself.
And I feel even the distinction between Chain and Model callbacks are not actually necessary if you have the Chain available in the in the callback, you can go up to the Model.

[brainlid]

Hi @tubedude! Thanks for taking the time to explain and document your use-case and reasoning.

I like the idea of passing the context through. I hadn't considered that. Here's the challenge I'm trying to resolve:

A ChatModel module's job is to handle the actual communication with the LLM. Specifically, it converts external concepts to and from something specific that LLM. That's where the deltas are actually received. Additionally, a ChatModel's call function supports making calls without using a chain. This also makes it easier to test in isolation.

That being said, I think a changed to the ChatModel.call function behaviour to add a context map
could be done easily enough. That would allow the LLMChain to pass it's content through to the ChatModel.

However, I also agree with your point about the confusion of the callback handlers as to what belongs to the ChatModel and what belongs to the LLMChain. A combined, full set of callbacks could be used and set on the LLMChain which could be passed through in the .call function too.

At that point, it makes me wonder about just passing in the full LLMChain to the ChatModel.call function. Circular references aside, it might be a smoother experience.

What do you think? Would that adequately address your needs? Anything else I'm missing or should consider?