18 gemini foundation technicalBlueprint2of3

defmodule DSPEx.Adapter do
  alias DSPEx.{Example, Prediction, Signature}

  @doc """
  The behaviour for all DSPEx Adapters.
  """
  @callback format(signature :: module(), demos :: list(Example.t()), inputs :: map()) :: list(map())
  @callback parse(signature :: module(), response :: map()) :: map()
end

defmodule DSPEx.Adapter.Chat do
  @behaviour DSPEx.Adapter

  # Main entry point.
  @impl DSPEx.Adapter
  def format(signature, demos, inputs) do
    system_message = build_system_message(signature)
    demo_messages = build_demo_messages(signature, demos)
    user_message = build_user_message(signature, inputs)

    # The final `messages` list structure
    [system_message | (demo_messages ++ [user_message])]
  end

  # This function parses the raw text completion back into structured fields.
  @impl DSPEx.Adapter
  def parse(signature, response) do
    # In Layer 2, this will be more complex. For now, we assume a simple case.
    # The response from OpenAI client is response["choices"][0]["message"]["content"]
    raw_text = get_in(response, ["choices", 0, "message", "content"])
    
    # Use Regex to parse fields like [[ ## answer ## ]]\n The answer is...
    output_fields = Signature.output_fields(signature)
    
    Enum.reduce(output_fields, %{}, fn field_name, acc ->
      # Regex to find the content for a specific field
      pattern = ~r/\[\[ ## #{field_name} ## \]\]\s*(.*)/s
      
      case Regex.run(pattern, raw_text, capture: :all_but_first) do
        [content] -> Map.put(acc, field_name, String.trim(content))
        nil -> acc
      end
    end)
  end

  # --- Private Helper Functions ---
  
  defp build_system_message(signature) do
    # Combines instructions, field descriptions, etc.
    %{role: "system", content: "..."}
  end

  defp build_demo_messages(signature, demos) do
    # Loops through demos, creating alternating user/assistant messages.
    Enum.flat_map(demos, fn demo ->
      [
        build_user_message(signature, demo.inputs),
        build_assistant_message(signature, demo.labels)
      ]
    end)
  end
  
  defp build_user_message(signature, inputs) do
    # Formats inputs into the "[[ ## field_name ## ]]\n..." format
    %{role: "user", content: "..."}
  end

  defp build_assistant_message(signature, outputs) do
    # Formats outputs into the "[[ ## field_name ## ]]\n..." format
    %{role: "assistant", content: "..."}
  end
end

defmodule DSPEx.ChainOfThought do
  @behaviour DSPEx.Program
  
  # The struct holds the original signature and an internal Predict module.
  defstruct [:signature, :predictor]

  # The constructor is where the magic happens.
  def new(signature, client) do
    # 1. Programmatically create a new signature with `reasoning` added.
    #    The `__using__` block of the signature macro will need to support
    #    being called with a map of fields, not just a string.
    original_fields = Signature.fields(signature)
    
    reasoning_field = {
      :reasoning, 
      %{desc: "Think step-by-step to reach the answer.", type: :string}
    }
    
    # Prepend the reasoning field before the original output fields.
    extended_fields = # logic to insert reasoning_field into original_fields
    extended_signature_name = :"#{inspect(signature)}_with_CoT"
    
    # Dynamically define the new signature module
    DynamicSupervisor.start_child(
        Application.get_env(:my_app, :dynamic_supervisor), 
        %{id: extended_signature_name, start: {Kernel, :make_module, [extended_signature_name, extended_fields]}}
    )

    # 2. Create an internal Predict module that uses this new, extended signature.
    predictor = %DSPEx.Predict{
      signature: extended_signature,
      client: client
    }
    
    %__MODULE__{signature: signature, predictor: predictor}
  end

  @impl DSPEx.Program
  def forward(program, inputs) do
    # 3. Delegate the call to the internal predictor.
    #    The result will be a Prediction struct containing the `reasoning` field.
    DSPEx.Program.forward(program.predictor, inputs)
  end
end

@impl DSPEx.Program
def forward(program, inputs) do
  # Assume an adapter is configured, e.g., in the program struct.
  # %DSPEx.Predict{signature: sig, client: c, adapter: DSPEx.Adapter.Chat}
  adapter = program.adapter

  # 1. Format the prompt using the adapter. Demos are empty for now.
  messages = adapter.format(program.signature, [], inputs)
  
  # 2. Call the LM client.
  case DSPEx.Client.LM.request(program.client, messages) do
    {:ok, lm_response} ->
      # 3. Parse the response using the adapter.
      parsed_fields = adapter.parse(program.signature, lm_response)
      
      # 4. Create the final prediction.
      prediction = DSPEx.Prediction.from_completions([parsed_fields], program.signature)
      {:ok, prediction}
      
    {:error, reason} ->
      {:error, reason}
  end
end