12 gemini

defmodule DSPEx.Predict do
  @moduledoc """
  A simple DSPEx.Program for single-turn predictions.

  This module orchestrates the process of formatting a request using an adapter,
  sending it to a client, and parsing the response.
  """

  @behaviour DSPEx.Program

  alias DSPEx.{Client, Prediction, Adapter.Chat}

  defstruct [
    :signature,
    :client,
    adapter: Chat, # The adapter module, defaults to Chat
    demos: []      # A list of DSPEx.Example structs for few-shot demonstrations
  ]

  @type t :: %__MODULE__{
    signature: module(),
    client: atom(),
    adapter: module(),
    demos: list(DSPEx.Example.t())
  }

  @impl DSPEx.Program
  def forward(%__MODULE__{} = program, inputs) do
    with {:ok, messages} <- program.adapter.format(program.signature, inputs, program.demos),
         {:ok, response_body} <- Client.request(program.client, messages),
         {:ok, parsed_outputs} <- program.adapter.parse(program.signature, response_body) do

      # On the happy path, construct the final, structured prediction object.
      prediction = %Prediction{
        inputs: inputs,
        outputs: parsed_outputs,
        raw_response: response_body
      }
      {:ok, prediction}
    else
      # Any error from the adapter or client pipeline steps is propagated.
      # The error will be a tuple like {:error, reason}, which we return directly.
      error -> error
    end
  end

  @impl DSPEx.Program
  def configure(%__MODULE__{} = program, config) when is_map(config) do
    # Returns a *new* struct with the updated values from the config map.
    # This is how optimizers will inject new demos into the program.
    # Example: configure(program, %{demos: new_demos_list})
    Map.merge(program, config)
  end
end

defmodule DSPEx.Evaluate do
  @moduledoc """
  A concurrent evaluation engine for DSPEx programs.
  """

  alias DSPEx.Program

  @default_opts [
    num_threads: System.schedulers_online() * 2,
    display_progress: true
  ]

  @doc """
  Runs a given program against a development set and computes a metric.

  Returns a tuple: `{:ok, aggregate_score, successes, failures}`.

  ## Options
    * `:num_threads` - The maximum number of concurrent processes. Defaults to
      twice the number of online schedulers.
    * `:display_progress` - Whether to show a progress bar. Defaults to `true`.
  """
  def run(program, devset, metric_fun, opts \\ []) do
    opts = Keyword.merge(@default_opts, opts)
    num_threads = opts[:num_threads]
    pbar = if opts[:display_progress], do: pbar(devset), else: nil

    results =
      devset
      |> Task.async_stream(&evaluate_example(&1, program, metric_fun),
        max_concurrency: num_threads,
        on_timeout: :kill_task
      )
      |> Stream.map(fn {:ok, result} ->
        if pbar, do: ProgressBar.tick(pbar)
        result
      end)
      |> Enum.to_list()

    if pbar, do: ProgressBar.done(pbar)

    process_results(results)
  end

  # This is the private function executed by each concurrent process.
  defp evaluate_example(example, program, metric_fun) do
    case Program.forward(program, example.inputs) do
      {:ok, prediction} ->
        # The metric function compares the original example (with labels)
        # to the prediction and returns a score (typically 0.0 to 1.0).
        score = metric_fun.(example, prediction)
        {:ok, {example, prediction, score}}

      {:error, reason} ->
        # A single failure does not stop the stream. It is returned as an error tuple.
        {:error, {example, reason}}
    end
  end

  # Processes the collected results from all tasks.
  defp process_results(results) do
    {successes, failures} = Enum.partition(results, fn
      {:ok, _} -> true
      {:error, _} -> false
    end)

    successes = Enum.map(successes, fn {:ok, res} -> res end)
    failures = Enum.map(failures, fn {:error, res} -> res end)

    # Calculate the aggregate score (e.g., average).
    total_score =
      successes
      |> Enum.map(fn {_ex, _pred, score} -> score end)
      |> Enum.sum()

    aggregate_score =
      if Enum.empty?(successes),
        do: 0.0,
        else: total_score / length(successes)

    {:ok, aggregate_score, successes, failures}
  end

  defp pbar(devset) do
    ProgressBar.new(
      total: Enum.count(devset),
      title: "Evaluating",
      format: ":title [:bar] :percent/:total | Rate: :rate/s | ETA: :eta"
    )
  end
end