01 gemini

defmodule DSPEx.Client do
  @moduledoc """
  A supervised GenServer for managing interactions with a specific AI provider.
  """
  use GenServer

  # --- Public API ---

  def start_link(opts) do
    provider_id = Keyword.fetch!(opts, :provider_id)
    # Register with foundation's service registry
    GenServer.start_link(__MODULE__, opts, name: via_tuple(provider_id))
  end

  def request(provider_id, messages, opts \\ %{}) do
    GenServer.call(via_tuple(provider_id), {:request, messages, opts})
  end

  # --- GenServer Callbacks ---

  @impl true
  def init(opts) do
    provider_id = Keyword.fetch!(opts, :provider_id)
    {:ok, config} = Foundation.Config.get([:dspex, :providers, provider_id])

    # Configure a dedicated connection pool and protection policy
    pool_name = :"#{provider_id}_pool"
    protection_key = :"#{provider_id}_protection"

    Foundation.Infrastructure.ConnectionManager.start_pool(pool_name,
      worker_module: DSPEx.Client.HttpWorker,
      worker_args: [config: config],
      size: config.pool_size || 10
    )

    Foundation.Infrastructure.configure_protection(protection_key, %{
      circuit_breaker: config.circuit_breaker,
      # Multi-bucket rate limiting for RPM and TPM
      rate_limiter: [
        rpm: {config.rate_limit_rpm.window, config.rate_limit_rpm.limit},
        tpm: {config.rate_limit_tpm.window, config.rate_limit_tpm.limit}
      ]
    })

    state = %{
      provider_id: provider_id,
      config: config,
      pool_name: pool_name,
      protection_key: protection_key
    }
    {:ok, state}
  end

  @impl true
  def handle_call({:request, messages, opts}, _from, state) do
    # Define the core API call function
    api_call_fun = fn worker_pid ->
      DSPEx.Client.HttpWorker.post(worker_pid, state.config.path, %{
        model: opts[:model] || state.config.default_model,
        messages: messages
        # ... other parameters
      })
    end

    # Calculate token cost for rate limiting
    token_cost = estimate_tokens(messages)

    # Use foundation's protection wrapper for the call
    response =
      Foundation.Infrastructure.execute_protected(
        state.protection_key,
        [
          connection_pool: state.pool_name,
          # Apply costs to the appropriate rate limit buckets
          rate_limiter_costs: [rpm: 1, tpm: token_cost]
        ],
        api_call_fun
      )

    # Emit telemetry via foundation
    emit_usage_telemetry(response, state, token_cost)
    {:reply, response, state}
  end

  # --- Private Helpers ---
  defp via_tuple(provider_id), do: Foundation.ServiceRegistry.via_tuple(:dspex, {:client, provider_id})
  defp estimate_tokens(messages), do: # ... token estimation logic ...
  defp emit_usage_telemetry(response, state, token_cost) do
    # ... logic to emit Foundation.Telemetry events for cost and latency ...
  end
end

defmodule DSPEx.PromptTemplate do
  @moduledoc "A composable engine for building and optimizing prompts."
  defstruct [:id, :name, :template_ast, :variables, :metadata, :parent]

  def create(name, template_string, opts \\ []) do
    # Parses the string into an AST, identifies variables
    %__MODULE__{id: Foundation.Utils.generate_id(), name: name, ...}
  end

  def render(template, variables) do
    # Renders the AST with variables, handling inheritance from parent templates
  end

  def compose(template1, template2) do
    # Creates a new template that combines the two
  end

  def optimize(template, performance_data) do
    # Suggests improvements based on historical data (e.g., adding few-shot examples)
  end
end

defmodule DSPEx.Cache do
  @moduledoc "A multi-tier cache with semantic retrieval capabilities."
  use GenServer

  # L1: ETS (in-memory) for exact matches (fastest)
  # L2: On-disk (persistent) for exact matches
  # L3: Vector DB (semantic) for similar matches

  defstruct [:l1_cache, :l2_cache, :vector_store]

  def get_or_compute(cache_key, prompt_text, compute_fun, opts \\ []) do
    # 1. Check L1 cache for exact match
    # 2. Check L2 cache for exact match
    # 3. If enabled, generate embedding for prompt_text and check L3 for similar prompts
    # 4. If all miss, execute `compute_fun`, then cache the result in L1, L2, and L3
  end

  # Helper module for generating embeddings for caching
  defmodule Embedding do
    def generate(text), do: # ... logic to call a local or API-based embedding model ...
  end
end

defmodule DSPEx.Analytics.CostManager do
  @moduledoc "Tracks and analyzes costs of LM operations."
  use GenServer

  def start_link(opts) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  def init(_opts) do
    # Subscribe to foundation events related to LM calls
    Foundation.Events.subscribe([:lm_request_complete])
    {:ok, %{cost_data: %{}}}
  end

  def handle_info({:foundation_event, event}, state) do
    # Process event, extract token counts and cost, update internal state
    # ...
    {:noreply, new_state}
  end

  def get_cost_analysis(time_range, group_by \\ [:provider, :model]) do
    # Query internal state to provide detailed cost breakdowns
  end
end

defmodule DSPEx.Workflow do
  @moduledoc "A declarative engine for running multi-step AI pipelines."

  @type step :: {step_name :: atom(), step_fun :: (map() -> {:ok, map()} | {:error, term()})}
  @type pipeline :: [step()]

  def run(pipeline, initial_data, opts \\ []) do
    # 1. Create a top-level Foundation.ErrorContext for the whole workflow.
    correlation_id = Keyword.get(opts, :correlation_id) || Foundation.Utils.generate_correlation_id()
    parent_context = Foundation.ErrorContext.new(__MODULE__, :run, correlation_id: correlation_id)

    Foundation.ErrorContext.with_context(parent_context, fn ->
      # 2. Use Enum.reduce_while to execute the pipeline.
      Enum.reduce_while(pipeline, {:ok, initial_data}, fn {step_name, fun}, {:ok, acc} ->
        # 3. For each step, create a child context and measure execution with Telemetry.
        step_context = Foundation.ErrorContext.child_context(parent_context, __MODULE__, step_name)

        {result, duration} = Foundation.Telemetry.measure([:dspex, :workflow, :step], %{}, fn ->
          Foundation.ErrorContext.with_context(step_context, fn -> fun.(acc) end)
        end)

        # 4. Handle step result and continue or halt.
        case result do
          {:ok, new_acc} -> {:cont, {:ok, Map.merge(acc, new_acc)}}
          {:error, reason} -> {:halt, {:error, reason}}
        end
      end)
    end)
  end
end

defmodule DSPEx.Predict do
  defstruct [:signature, :client_id, :template]

  def forward(%__MODULE__{} = program, inputs) do
    # 1. Define the workflow using the DSPex Foundation engine.
    pipeline = [
      {:render_prompt, &render_prompt_step/1},
      {:call_lm, &call_lm_step/1},
      {:parse_response, &parse_response_step/1}
    ]

    # 2. The initial data for the workflow includes the program state and inputs.
    initial_data = %{
      program: program,
      inputs: inputs
    }

    # 3. Run the workflow.
    DSPEx.Workflow.run(pipeline, initial_data)
  end

  # --- Workflow Step Implementations ---

  defp render_prompt_step(%{program: prog, inputs: inputs}) do
    # Uses the PromptTemplate engine
    rendered_prompt = DSPEx.PromptTemplate.render(prog.template, Map.merge(inputs, prog.signature))
    {:ok, %{prompt_messages: [role: "user", content: rendered_prompt]}}
  end

  defp call_lm_step(%{program: prog, prompt_messages: messages}) do
    # Uses the AI Provider Client
    case DSPEx.Client.request(prog.client_id, messages) do
      {:ok, response} -> {:ok, %{lm_response: response}}
      error -> error
    end
  end

  defp parse_response_step(%{program: prog, lm_response: response}) do
    # Logic to parse the response based on prog.signature
    # For simplicity, we just extract the content here.
    answer = get_in(response, [:choices, 0, :message, :content])
    {:ok, %{answer: answer}}
  end
end