|
| 1 | +# SPDX-FileCopyrightText: Sudo Apt Holdings LLC |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +defmodule Trinity.LLM do |
| 4 | + @moduledoc """ |
| 5 | + The one door to a language model. Slice 011. |
| 6 | +
|
| 7 | + Sessions call this and never a provider. A request names a registry model id (or nothing, for |
| 8 | + the default); the registry names the provider module; the call is retried on transient errors; |
| 9 | + a completed call writes a `usage_events` row with cost from the registry price. Events reach |
| 10 | + the caller either through a function (`stream/3`) or as messages to a pid (`stream_to/3`), |
| 11 | + which is what slice 012's Session wants. |
| 12 | + """ |
| 13 | + use Boundary, deps: [Trinity], exports: [Error, Event, Request, Provider, Registry] |
| 14 | + |
| 15 | + alias Trinity.LLM.{Error, Registry, Request, Retry, Usage} |
| 16 | + |
| 17 | + @type opts :: keyword() |
| 18 | + |
| 19 | + @doc "Streams events through `emit`, in the calling process. Returns the usage." |
| 20 | + @spec stream(Request.t(), opts(), Trinity.LLM.Provider.emit()) :: |
| 21 | + {:ok, map()} | {:error, Error.t() | term()} |
| 22 | + def stream(%Request{} = request, opts \\ [], emit) when is_function(emit, 1) do |
| 23 | + with {:ok, entry, module} <- resolve(request) do |
| 24 | + call(entry, "chat", opts, fn -> |
| 25 | + module.stream(request, provider_opts(entry, opts), emit) |
| 26 | + end) |
| 27 | + end |
| 28 | + end |
| 29 | + |
| 30 | + @doc """ |
| 31 | + Streams events as messages `{:llm_event, ref, event}` to `pid`, then `{:llm_done, ref, result}`. |
| 32 | + Runs in a supervised task so the caller is never blocked; `ref` is returned at once. |
| 33 | + """ |
| 34 | + @spec stream_to(Request.t(), opts(), pid()) :: {:ok, reference()} | {:error, term()} |
| 35 | + def stream_to(%Request{} = request, opts \\ [], pid) when is_pid(pid) do |
| 36 | + ref = make_ref() |
| 37 | + |
| 38 | + with {:ok, _entry, _module} <- resolve(request), |
| 39 | + {:ok, _task} <- |
| 40 | + Task.Supervisor.start_child(Trinity.LLM.TaskSupervisor, fn -> |
| 41 | + result = stream(request, opts, &send(pid, {:llm_event, ref, &1})) |
| 42 | + send(pid, {:llm_done, ref, result}) |
| 43 | + end) do |
| 44 | + {:ok, ref} |
| 45 | + end |
| 46 | + end |
| 47 | + |
| 48 | + @doc "One complete response." |
| 49 | + @spec generate(Request.t(), opts()) :: {:ok, Trinity.LLM.Provider.result()} | {:error, term()} |
| 50 | + def generate(%Request{} = request, opts \\ []) do |
| 51 | + with {:ok, entry, module} <- resolve(request) do |
| 52 | + call(entry, "chat", opts, fn -> module.generate(request, provider_opts(entry, opts)) end) |
| 53 | + end |
| 54 | + end |
| 55 | + |
| 56 | + @doc "A map validated against a JSON Schema." |
| 57 | + @spec generate_object(Request.t(), map(), opts()) :: {:ok, map()} | {:error, term()} |
| 58 | + def generate_object(%Request{} = request, schema, opts \\ []) when is_map(schema) do |
| 59 | + with {:ok, entry, module} <- resolve(request), |
| 60 | + {:ok, object, _usage} <- |
| 61 | + call(entry, "object", opts, fn -> |
| 62 | + module.generate_object(request, schema, provider_opts(entry, opts)) |
| 63 | + end) do |
| 64 | + {:ok, object} |
| 65 | + end |
| 66 | + end |
| 67 | + |
| 68 | + @doc "One vector per text. `opts[:model]` names the embedding model's registry id." |
| 69 | + @spec embed([String.t()], opts()) :: {:ok, [[float()]]} | {:error, term()} |
| 70 | + def embed(texts, opts \\ []) when is_list(texts) do |
| 71 | + with {:ok, entry} <- Registry.lookup(Keyword.get(opts, :model)), |
| 72 | + {:ok, module} <- Registry.provider_module(entry), |
| 73 | + {:ok, vectors, _usage} <- |
| 74 | + call(entry, "embed", opts, fn -> module.embed(texts, provider_opts(entry, opts)) end) do |
| 75 | + {:ok, vectors} |
| 76 | + end |
| 77 | + end |
| 78 | + |
| 79 | + @doc "The registry's models." |
| 80 | + @spec models() :: [Registry.entry()] |
| 81 | + def models, do: Registry.models() |
| 82 | + |
| 83 | + @doc "The registry's default model id." |
| 84 | + @spec default_model() :: String.t() | nil |
| 85 | + def default_model, do: Registry.default_model() |
| 86 | + |
| 87 | + @doc "A model's capabilities, from its registry entry." |
| 88 | + @spec capabilities(String.t()) :: {:ok, [atom() | {atom(), term()}]} | {:error, term()} |
| 89 | + def capabilities(model_id) do |
| 90 | + with {:ok, entry} <- Registry.lookup(model_id), do: {:ok, entry.caps} |
| 91 | + end |
| 92 | + |
| 93 | + defp resolve(%Request{model: model}) do |
| 94 | + with {:ok, entry} <- Registry.lookup(model), |
| 95 | + {:ok, module} <- Registry.provider_module(entry) do |
| 96 | + {:ok, entry, module} |
| 97 | + end |
| 98 | + end |
| 99 | + |
| 100 | + # The entry's keys win: a caller's `:model` is a registry id, the entry's is the provider's |
| 101 | + # own name, and the provider must see the latter. Found by the live suite, where an embed |
| 102 | + # call asked req_llm for a provider named after the registry id. |
| 103 | + defp provider_opts(entry, opts) do |
| 104 | + entry_opts = |
| 105 | + entry |> Map.take([:model, :base_url, :api_key_env, :price, :caps]) |> Map.to_list() |
| 106 | + |
| 107 | + Keyword.merge(opts, entry_opts) |
| 108 | + end |
| 109 | + |
| 110 | + # Retry around the provider; on success the usage is recorded. The usage row is written for |
| 111 | + # the call that completed, once, whatever the number of attempts it took. |
| 112 | + defp call(entry, kind, opts, fun) do |
| 113 | + case Retry.run(fun, opts) do |
| 114 | + {:ok, %{usage: usage}} = ok -> |
| 115 | + record(entry, kind, usage, opts) |
| 116 | + ok |
| 117 | + |
| 118 | + {:ok, usage} = ok when is_map(usage) -> |
| 119 | + record(entry, kind, usage, opts) |
| 120 | + ok |
| 121 | + |
| 122 | + {:ok, _value, usage} = ok -> |
| 123 | + record(entry, kind, usage, opts) |
| 124 | + ok |
| 125 | + |
| 126 | + other -> |
| 127 | + other |
| 128 | + end |
| 129 | + end |
| 130 | + |
| 131 | + defp record(entry, kind, usage, opts) do |
| 132 | + {:ok, _} = Usage.record(entry, kind, usage, Keyword.take(opts, [:session_id])) |
| 133 | + end |
| 134 | +end |
0 commit comments