252 lines
7.1 KiB
Elixir
252 lines
7.1 KiB
Elixir
defmodule UncloseAI do
|
|
@moduledoc """
|
|
UncloseAI Elixir Library
|
|
OpenAI-compatible API client with streaming support
|
|
Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
|
|
"""
|
|
|
|
defmodule ModelInfo do
|
|
@moduledoc "Struct to hold model information"
|
|
defstruct [:id, :endpoint, :max_tokens]
|
|
end
|
|
|
|
defmodule Client do
|
|
@moduledoc "Client struct to hold discovered models and endpoints"
|
|
defstruct models: [], tts_endpoints: [], timeout: 30_000
|
|
|
|
@doc """
|
|
Initialize a new UncloseAI client with auto-discovery
|
|
"""
|
|
def new(opts \\ []) do
|
|
timeout = Keyword.get(opts, :timeout, 30_000)
|
|
IO.puts("Initializing UncloseAI client...")
|
|
|
|
# Discover chat/code models
|
|
models =
|
|
Stream.iterate(1, &(&1 + 1))
|
|
|> Stream.take_while(fn i ->
|
|
System.get_env("MODEL_ENDPOINT_#{i}") != nil
|
|
end)
|
|
|> Stream.map(fn i ->
|
|
endpoint = System.get_env("MODEL_ENDPOINT_#{i}")
|
|
IO.puts("Endpoint #{i}: #{endpoint}")
|
|
discover_models_from_endpoint(endpoint)
|
|
end)
|
|
|> Enum.to_list()
|
|
|> List.flatten()
|
|
|
|
# Discover TTS endpoints
|
|
tts_endpoints =
|
|
Stream.iterate(1, &(&1 + 1))
|
|
|> Stream.take_while(fn i ->
|
|
System.get_env("TTS_ENDPOINT_#{i}") != nil
|
|
end)
|
|
|> Stream.map(fn i ->
|
|
System.get_env("TTS_ENDPOINT_#{i}")
|
|
end)
|
|
|> Enum.to_list()
|
|
|
|
IO.puts("Discovered #{length(models)} models, #{length(tts_endpoints)} TTS endpoints\n")
|
|
|
|
%Client{models: models, tts_endpoints: tts_endpoints, timeout: timeout}
|
|
end
|
|
|
|
defp discover_models_from_endpoint(endpoint) do
|
|
case HTTPoison.get("#{endpoint}/models", [], recv_timeout: 10_000) do
|
|
{:ok, %HTTPoison.Response{body: body}} ->
|
|
data = Jason.decode!(body)
|
|
|
|
case data["data"] do
|
|
nil ->
|
|
[]
|
|
|
|
models ->
|
|
Enum.filter(models, fn model ->
|
|
# Skip modelperm-* entries
|
|
not String.starts_with?(model["id"], "modelperm-")
|
|
end)
|
|
|> Enum.map(fn model ->
|
|
model_id = model["id"]
|
|
max_tokens = model["max_model_len"] || 8192
|
|
%ModelInfo{id: model_id, endpoint: endpoint, max_tokens: max_tokens}
|
|
end)
|
|
end
|
|
|
|
{:error, _reason} ->
|
|
# Silently skip failed endpoints
|
|
[]
|
|
end
|
|
end
|
|
|
|
@doc """
|
|
Non-streaming chat completion
|
|
"""
|
|
def chat(client, messages, opts \\ []) do
|
|
model_idx = Keyword.get(opts, :model_idx, 0)
|
|
max_tokens = Keyword.get(opts, :max_tokens, 100)
|
|
temperature = Keyword.get(opts, :temperature, 0.7)
|
|
|
|
model = Enum.at(client.models, model_idx)
|
|
|
|
if model == nil do
|
|
{:error, "Invalid model index"}
|
|
else
|
|
url = "#{model.endpoint}/chat/completions"
|
|
|
|
request = %{
|
|
model: model.id,
|
|
messages: messages,
|
|
stream: false,
|
|
max_tokens: max_tokens,
|
|
temperature: temperature
|
|
}
|
|
|
|
case HTTPoison.post(
|
|
url,
|
|
Jason.encode!(request),
|
|
[{"Content-Type", "application/json"}],
|
|
recv_timeout: client.timeout
|
|
) do
|
|
{:ok, %HTTPoison.Response{body: body}} ->
|
|
data = Jason.decode!(body)
|
|
{:ok, get_in(data, ["choices", Access.at(0), "message", "content"])}
|
|
|
|
{:error, reason} ->
|
|
{:error, reason}
|
|
end
|
|
end
|
|
end
|
|
|
|
@doc """
|
|
Streaming chat completion - returns a Stream that yields content chunks
|
|
"""
|
|
def chat_stream(client, messages, opts \\ []) do
|
|
model_idx = Keyword.get(opts, :model_idx, 0)
|
|
max_tokens = Keyword.get(opts, :max_tokens, 500)
|
|
temperature = Keyword.get(opts, :temperature, 0.7)
|
|
|
|
model = Enum.at(client.models, model_idx)
|
|
|
|
if model == nil do
|
|
raise "Invalid model index"
|
|
end
|
|
|
|
url = "#{model.endpoint}/chat/completions"
|
|
|
|
request = %{
|
|
model: model.id,
|
|
messages: messages,
|
|
stream: true,
|
|
max_tokens: max_tokens,
|
|
temperature: temperature
|
|
}
|
|
|
|
# Use HTTPoison stream with async response handling
|
|
Stream.resource(
|
|
fn ->
|
|
{:ok, response} =
|
|
HTTPoison.post(
|
|
url,
|
|
Jason.encode!(request),
|
|
[{"Content-Type", "application/json"}],
|
|
stream_to: self(),
|
|
async: :once,
|
|
recv_timeout: client.timeout
|
|
)
|
|
|
|
{response, ""}
|
|
end,
|
|
fn {response, buffer} ->
|
|
receive do
|
|
%HTTPoison.AsyncStatus{} ->
|
|
HTTPoison.stream_next(response)
|
|
{[], {response, buffer}}
|
|
|
|
%HTTPoison.AsyncHeaders{} ->
|
|
HTTPoison.stream_next(response)
|
|
{[], {response, buffer}}
|
|
|
|
%HTTPoison.AsyncChunk{chunk: chunk} ->
|
|
# Append chunk to buffer and process lines
|
|
new_buffer = buffer <> chunk
|
|
{lines, remaining} = extract_lines(new_buffer)
|
|
|
|
contents =
|
|
lines
|
|
|> Enum.filter(&String.starts_with?(&1, "data: "))
|
|
|> Enum.map(&String.replace_prefix(&1, "data: ", ""))
|
|
|> Enum.reject(&(&1 == "[DONE]"))
|
|
|> Enum.map(&extract_content/1)
|
|
|> Enum.reject(&is_nil/1)
|
|
|
|
HTTPoison.stream_next(response)
|
|
{contents, {response, remaining}}
|
|
|
|
%HTTPoison.AsyncEnd{} ->
|
|
{:halt, {response, buffer}}
|
|
after
|
|
client.timeout ->
|
|
{:halt, {response, buffer}}
|
|
end
|
|
end,
|
|
fn {response, _buffer} ->
|
|
:hackney.close(response.id)
|
|
end
|
|
)
|
|
end
|
|
|
|
defp extract_lines(buffer) do
|
|
lines = String.split(buffer, "\n")
|
|
|
|
case List.last(lines) do
|
|
"" -> {Enum.drop(lines, -1), ""}
|
|
partial -> {Enum.drop(lines, -1), partial}
|
|
end
|
|
end
|
|
|
|
defp extract_content(data) do
|
|
case Jason.decode(data) do
|
|
{:ok, json} ->
|
|
get_in(json, ["choices", Access.at(0), "delta", "content"])
|
|
|
|
{:error, _} ->
|
|
nil
|
|
end
|
|
end
|
|
|
|
@doc """
|
|
Text-to-speech generation
|
|
"""
|
|
def tts(client, text, opts \\ []) do
|
|
voice = Keyword.get(opts, :voice, "alloy")
|
|
output_file = Keyword.get(opts, :output_file, "/tmp/speech.mp3")
|
|
|
|
if Enum.empty?(client.tts_endpoints) do
|
|
{:error, "No TTS endpoints available"}
|
|
else
|
|
endpoint = List.first(client.tts_endpoints)
|
|
url = "#{endpoint}/audio/speech"
|
|
|
|
request = %{
|
|
model: "tts-1",
|
|
voice: voice,
|
|
input: text
|
|
}
|
|
|
|
case HTTPoison.post(
|
|
url,
|
|
Jason.encode!(request),
|
|
[{"Content-Type", "application/json"}],
|
|
recv_timeout: client.timeout
|
|
) do
|
|
{:ok, %HTTPoison.Response{body: body}} ->
|
|
File.write!(output_file, body)
|
|
{:ok, output_file}
|
|
|
|
{:error, reason} ->
|
|
{:error, reason}
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end
|