langchain

defmodule LangChain.Chains.ChainCallbacks do
  @moduledoc """
  Defines the callbacks fired by an LLMChain and LLM module.

  A callback handler is a map that defines the specific callback event with a
  function to execute for that event.

  ## Example

  A sample configured callback handler that forwards received data to a specific
  LiveView.

      live_view_pid = self()

      my_handlers = %{
        on_llm_new_delta: fn _chain, new_deltas -> send(live_view_pid, {:received_delta, new_deltas}) end,
        on_message_processed: fn _chain, new_message -> send(live_view_pid, {:received_message, new_message}) end,
        on_error_message_created: fn _chain, new_message -> send(live_view_pid, {:received_message, new_message}) end
      }

      model = SomeLLM.new!(%{...})

      chain =
        %{llm: model}
        |> LLMChain.new!()
        |> LLMChain.add_callback(my_handlers)

  """

  alias LangChain.Chains.LLMChain
  alias LangChain.Function
  alias LangChain.Message
  alias LangChain.Message.ToolCall
  alias LangChain.Message.ToolResult
  alias LangChain.MessageDelta
  alias LangChain.TokenUsage

  @typedoc """
  Executed when an LLM is streaming a response and a new MessageDelta (or token)
  was received.

  - `:index` is optionally present if the LLM supports sending `n` versions of a
    response.

  The return value is discarded.
  """
  @type llm_new_delta :: (LLMChain.t(), [MessageDelta.t()] -> any())

  @typedoc """
  Executed when an LLM is not streaming and a full message was received.

  The return value is discarded.
  """
  @type llm_new_message :: (LLMChain.t(), Message.t() -> any())

  @typedoc """
  Executed when an LLM (typically a service) responds with rate limiting
  information.

  The specific rate limit information depends on the LLM. It returns a map with
  all the available information included.

  The return value is discarded.
  """
  @type llm_ratelimit_info :: (LLMChain.t(), info :: %{String.t() => any()} -> any())

  @typedoc """
  Executed when an LLM response reports the token usage in a
  `LangChain.TokenUsage` struct. The data returned depends on the LLM.

  The return value is discarded.
  """
  @type llm_token_usage :: (LLMChain.t(), TokenUsage.t() -> any())

  @typedoc """
  Executed when an LLM response is received through an HTTP response. The entire
  set of raw response headers can be received and processed.

  The return value is discarded.

  ## Example

  A function declaration that matches the signature.

      def handle_llm_response_headers(chain, response_headers) do
        # This demonstrates how to send the response headers to a
        # LiveView assuming the LiveView's pid was stored in the chain's
        # custom_context.
        send(chain.custom_context.live_view_pid, {:req_response_headers, response_headers})

        IO.inspect(response_headers)
      end
  """
  @type llm_response_headers :: (LLMChain.t(), response_headers :: map() -> any())

  @typedoc """
  Executed when an LLMChain has completed processing a received assistant
  message. This fires when a message is complete either after assembling
  streaming deltas or when a full message is received when not streaming.

  This is the best way to be notified when a message is "done" and should be
  handled by the application.

  The handler's return value is discarded.
  """
  @type chain_message_processed :: (LLMChain.t(), Message.t() -> any())

  @typedoc """
  Executed when an LLMChain, in response to an error from the LLM, generates a
  new, automated response message intended to be returned to the LLM.

  """
  @type chain_error_message_created :: (LLMChain.t(), Message.t() -> any())

  @typedoc """
  Executed when processing a received message errors or fails. The erroring
  message is included in the callback with the state of processing that was
  completed before erroring.

  The handler's return value is discarded.
  """
  @type chain_message_processing_error :: (LLMChain.t(), Message.t() -> any())

  @typedoc """
  Executed when a tool call is identified during streaming, before execution begins.

  This fires as soon as we have enough information to identify the tool (at minimum, the `name` field).
  The tool call may be incomplete - `call_id` might not be available yet, and `arguments` may be partial.

  This callback provides early notification for UI feedback like "Searching web..." while the LLM
  is still streaming the complete tool call.

  Timing:
  - Fires: As soon as tool name is detected in streaming deltas
  - Before: Tool arguments are fully received
  - Before: Tool execution begins

  Arguments:
  - First: LLMChain.t() - Current chain state
  - Second: ToolCall.t() - Tool call struct (may be incomplete, but has name)
  - Third: Function.t() - Function definition (includes display_text)

  The handler's return value is discarded.

  ## Example

      callback_handler = %{
        on_tool_call_identified: fn _chain, tool_call, func ->
          IO.puts("Tool identified: \#{func.display_text || tool_call.name}")
        end
      }

  """
  @type chain_tool_call_identified :: (LLMChain.t(), ToolCall.t(), Function.t() -> any())

  @typedoc """
  Executed when the chain begins executing a tool call.

  This fires immediately before tool execution starts, allowing UIs to show
  real-time feedback like "Searching the web..." or "Creating file...".

  - First argument: LLMChain.t()
  - Second argument: ToolCall struct being executed
  - Third argument: Function struct for the tool (includes display_text)

  The handler's return value is discarded.
  """
  @type chain_tool_execution_started :: (LLMChain.t(), ToolCall.t(), Function.t() -> any())

  @typedoc """
  Executed when a single tool execution completes successfully.

  Fires after individual tool execution, before results are aggregated.
  Useful for showing per-tool success indicators.

  - First argument: LLMChain.t()
  - Second argument: ToolCall that was executed
  - Third argument: ToolResult that was generated

  The handler's return value is discarded.
  """
  @type chain_tool_execution_completed :: (LLMChain.t(), ToolCall.t(), ToolResult.t() -> any())

  @typedoc """
  Executed when a single tool execution fails.

  Fires when tool execution raises an exception or returns an error result.

  - First argument: LLMChain.t()
  - Second argument: ToolCall that failed
  - Third argument: Error reason or exception

  The handler's return value is discarded.
  """
  @type chain_tool_execution_failed :: (LLMChain.t(), ToolCall.t(), term() -> any())

  @typedoc """
  Executed when one or more tools return an interrupt signal.

  Fires once per tool execution batch with all interrupted results.
  The tool is paused and awaiting external input to continue.

  - First argument: LLMChain.t()
  - Second argument: List of ToolResult structs with `is_interrupt: true`

  The handler's return value is discarded.
  """
  @type chain_tool_interrupted :: (LLMChain.t(), [ToolResult.t()] -> any())

  @typedoc """
  Executed when the chain uses one or more tools and the resulting ToolResults
  are generated as part of a tool response message.

  The handler's return value is discarded.
  """
  @type chain_tool_response_created :: (LLMChain.t(), Message.t() -> any())

  @typedoc """
  Executed when the chain failed multiple times used up the `max_retry_count`
  resulting in the process aborting and returning an error.

  The handler's return value is discarded.
  """
  @type chain_retries_exceeded :: (LLMChain.t() -> any())

  @typedoc """
  The supported set of callbacks for an LLM module.
  """
  @type chain_callback_handler :: %{
          # model-level callbacks
          optional(:on_llm_new_delta) => llm_new_delta(),
          optional(:on_llm_new_message) => llm_new_message(),
          optional(:on_llm_ratelimit_info) => llm_ratelimit_info(),
          optional(:on_llm_token_usage) => llm_token_usage(),
          optional(:on_llm_response_headers) => llm_response_headers(),

          # Chain-level callbacks
          optional(:on_message_processed) => chain_message_processed(),
          optional(:on_message_processing_error) => chain_message_processing_error(),
          optional(:on_error_message_created) => chain_error_message_created(),
          optional(:on_tool_call_identified) => chain_tool_call_identified(),
          optional(:on_tool_execution_started) => chain_tool_execution_started(),
          optional(:on_tool_execution_completed) => chain_tool_execution_completed(),
          optional(:on_tool_execution_failed) => chain_tool_execution_failed(),
          optional(:on_tool_interrupted) => chain_tool_interrupted(),
          optional(:on_tool_response_created) => chain_tool_response_created(),
          optional(:on_retries_exceeded) => chain_retries_exceeded()
        }
end