934 lines
34 KiB
C#
934 lines
34 KiB
C#
using System.Collections.Concurrent;
|
|
using System.Net.WebSockets;
|
|
using System.Text;
|
|
using System.Text.Json;
|
|
using Microsoft.Extensions.Logging;
|
|
using Microsoft.Extensions.Logging.Abstractions;
|
|
using Ngino.Protocol;
|
|
|
|
namespace Ngino.Client;
|
|
|
|
internal sealed class TunnelClient
|
|
{
|
|
private static readonly TimeSpan ModelRefreshInterval = TimeSpan.FromSeconds(15);
|
|
private static readonly TimeSpan ModelRefreshTimeout = TimeSpan.FromSeconds(10);
|
|
private static readonly JsonSerializerOptions JsonOptions = new(JsonSerializerDefaults.Web);
|
|
private const string EmbeddingWarmupInput = "Ngino warmup";
|
|
|
|
private readonly ConcurrentDictionary<string, UpstreamRequest> _activeRequests = new();
|
|
private readonly ConcurrentDictionary<string, PendingRequestBody> _pendingRequestBodies = new();
|
|
private readonly HttpClient _httpClient;
|
|
private readonly ClientOptions _options;
|
|
private readonly ILogger<TunnelClient> _logger;
|
|
private readonly object _modelSnapshotLock = new();
|
|
private readonly SemaphoreSlim _sendLock = new(1, 1);
|
|
private readonly LlamaCppManager? _llamaCppManager;
|
|
private List<string> _lastActiveModels = [];
|
|
private List<string> _lastModels = [];
|
|
|
|
public TunnelClient(ClientOptions options, ILogger<TunnelClient>? logger = null)
|
|
{
|
|
_options = options;
|
|
_logger = logger ?? NullLogger<TunnelClient>.Instance;
|
|
_httpClient = new HttpClient
|
|
{
|
|
Timeout = Timeout.InfiniteTimeSpan
|
|
};
|
|
|
|
if (_options.UseLlamaCppViaDocker)
|
|
{
|
|
if (string.IsNullOrWhiteSpace(_options.UseOllamaModelsPath))
|
|
{
|
|
throw new InvalidOperationException(
|
|
"--use-ollama-models-path is required when --use-llama-cpp-via-docker is set.");
|
|
}
|
|
|
|
_llamaCppManager = new LlamaCppManager(
|
|
_options.UseOllamaModelsPath,
|
|
_options.LlamaCppDockerImage,
|
|
_options.LlamaCppBasePort,
|
|
_logger);
|
|
}
|
|
}
|
|
|
|
public async Task RunAsync(CancellationToken cancellationToken)
|
|
{
|
|
if (_llamaCppManager is not null)
|
|
{
|
|
_logger.LogInformation("Testing Docker availability...");
|
|
var dockerAvailable = await _llamaCppManager.TestDockerAsync();
|
|
if (!dockerAvailable)
|
|
{
|
|
_logger.LogWarning("Docker is not available. llama.cpp via Docker will not work.");
|
|
}
|
|
else
|
|
{
|
|
_logger.LogInformation("Docker is available. Using llama.cpp image: {Image}", _llamaCppManager.DockerImage);
|
|
}
|
|
}
|
|
|
|
while (!cancellationToken.IsCancellationRequested)
|
|
{
|
|
using var socket = new ClientWebSocket();
|
|
socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(30);
|
|
|
|
if (_options.InsecureSkipTlsVerify)
|
|
{
|
|
socket.Options.RemoteCertificateValidationCallback = static (_, _, _, _) => true;
|
|
}
|
|
|
|
if (!string.IsNullOrWhiteSpace(_options.Token))
|
|
{
|
|
socket.Options.SetRequestHeader(ProtocolConstants.TokenHeader, _options.Token);
|
|
}
|
|
|
|
socket.Options.SetRequestHeader(ProtocolConstants.ClientIdHeader, _options.ClientId);
|
|
|
|
using var connectionCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
|
Task? modelRefreshTask = null;
|
|
|
|
try
|
|
{
|
|
_logger.LogInformation("Connecting to {TunnelUri}...", _options.TunnelUri);
|
|
await socket.ConnectAsync(_options.TunnelUri, cancellationToken);
|
|
_logger.LogInformation("Tunnel connected.");
|
|
|
|
modelRefreshTask = RefreshModelsLoopAsync(socket, connectionCts.Token);
|
|
await ReceiveLoopAsync(socket, connectionCts.Token);
|
|
|
|
if (socket.CloseStatusDescription == ProtocolConstants.ReplacedCloseDescription)
|
|
{
|
|
_logger.LogWarning("This client was replaced by a newer tunnel client. Exiting.");
|
|
return;
|
|
}
|
|
|
|
_logger.LogInformation("Tunnel closed by server ({Reason}).", socket.CloseStatusDescription ?? "no reason given");
|
|
}
|
|
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
break;
|
|
}
|
|
catch (Exception exception)
|
|
{
|
|
_logger.LogWarning("Tunnel disconnected: {Message}", exception.Message);
|
|
}
|
|
finally
|
|
{
|
|
connectionCts.Cancel();
|
|
if (modelRefreshTask is not null)
|
|
{
|
|
try
|
|
{
|
|
await modelRefreshTask;
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
}
|
|
}
|
|
|
|
CancelAllActiveRequests();
|
|
}
|
|
|
|
if (!cancellationToken.IsCancellationRequested)
|
|
{
|
|
_logger.LogInformation("Reconnecting in {Seconds:0.#} seconds...", _options.ReconnectDelay.TotalSeconds);
|
|
await Task.Delay(_options.ReconnectDelay, cancellationToken);
|
|
}
|
|
}
|
|
}
|
|
|
|
private async Task RefreshModelsLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
|
|
{
|
|
while (socket.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
|
|
{
|
|
try
|
|
{
|
|
await RefreshModelsOnceAsync(socket, cancellationToken);
|
|
}
|
|
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
break;
|
|
}
|
|
catch (Exception exception)
|
|
{
|
|
_logger.LogWarning("Failed to report upstream model list: {Message}", exception.Message);
|
|
}
|
|
|
|
await Task.Delay(ModelRefreshInterval, cancellationToken);
|
|
}
|
|
}
|
|
|
|
private async Task RefreshModelsOnceAsync(ClientWebSocket socket, CancellationToken cancellationToken)
|
|
{
|
|
var modelsTask = TryRefreshModelListAsync(GetUpstreamModelsAsync, "listed", cancellationToken);
|
|
var activeModelsTask = TryRefreshModelListAsync(GetActiveUpstreamModelsAsync, "active", cancellationToken);
|
|
|
|
await Task.WhenAll(modelsTask, activeModelsTask);
|
|
|
|
var snapshot = UpdateCachedModelSnapshot(modelsTask.Result, activeModelsTask.Result);
|
|
|
|
await SendAsync(
|
|
socket,
|
|
new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelSnapshot,
|
|
Models = snapshot.Models,
|
|
ActiveModels = snapshot.ActiveModels
|
|
},
|
|
cancellationToken);
|
|
|
|
_logger.LogInformation(
|
|
"Reported {ModelCount} listed and {ActiveModelCount} active upstream model(s).",
|
|
snapshot.Models.Count,
|
|
snapshot.ActiveModels.Count);
|
|
}
|
|
|
|
private async Task<List<string>> GetUpstreamModelsAsync(CancellationToken cancellationToken)
|
|
{
|
|
if (_llamaCppManager is not null)
|
|
{
|
|
var models = _llamaCppManager.DiscoverModelsWithBlob();
|
|
return models
|
|
.Select(m => m.OllamaName)
|
|
.Where(name => !string.IsNullOrWhiteSpace(name))
|
|
.OrderBy(name => name, StringComparer.OrdinalIgnoreCase)
|
|
.ToList();
|
|
}
|
|
|
|
using var request = new HttpRequestMessage(HttpMethod.Get, new Uri(_options.Upstream, "/api/tags"));
|
|
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
|
response.EnsureSuccessStatusCode();
|
|
|
|
await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
|
|
using var document = await JsonDocument.ParseAsync(stream, cancellationToken: cancellationToken);
|
|
|
|
return ExtractModelNames(document.RootElement);
|
|
}
|
|
|
|
private async Task<List<string>> GetActiveUpstreamModelsAsync(CancellationToken cancellationToken)
|
|
{
|
|
if (_llamaCppManager is not null)
|
|
{
|
|
var models = _llamaCppManager.DiscoverModelsWithBlob();
|
|
return models
|
|
.Where(m => _llamaCppManager.IsModelActive(m.OllamaName))
|
|
.Select(m => m.OllamaName)
|
|
.Where(name => !string.IsNullOrWhiteSpace(name))
|
|
.OrderBy(name => name, StringComparer.OrdinalIgnoreCase)
|
|
.ToList();
|
|
}
|
|
|
|
using var request = new HttpRequestMessage(HttpMethod.Get, new Uri(_options.Upstream, "/api/ps"));
|
|
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
|
|
|
if (!response.IsSuccessStatusCode)
|
|
{
|
|
_logger.LogDebug("Ollama /api/ps returned {StatusCode}; active model list will be empty.", response.StatusCode);
|
|
return [];
|
|
}
|
|
|
|
await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
|
|
using var document = await JsonDocument.ParseAsync(stream, cancellationToken: cancellationToken);
|
|
|
|
return ExtractModelNames(document.RootElement);
|
|
}
|
|
|
|
internal static List<string> ExtractModelNames(JsonElement root)
|
|
{
|
|
var models = new List<string>();
|
|
|
|
if (root.ValueKind == JsonValueKind.Object
|
|
&& root.TryGetProperty("models", out var ollamaModels)
|
|
&& ollamaModels.ValueKind == JsonValueKind.Array)
|
|
{
|
|
AddModelNames(models, ollamaModels, "name");
|
|
AddModelNames(models, ollamaModels, "model");
|
|
}
|
|
|
|
return NormalizeModelNames(models);
|
|
}
|
|
|
|
private static void AddModelNames(List<string> models, JsonElement array, string propertyName)
|
|
{
|
|
foreach (var item in array.EnumerateArray())
|
|
{
|
|
if (item.ValueKind == JsonValueKind.Object
|
|
&& item.TryGetProperty(propertyName, out var model)
|
|
&& model.ValueKind == JsonValueKind.String
|
|
&& !string.IsNullOrWhiteSpace(model.GetString()))
|
|
{
|
|
models.Add(model.GetString()!);
|
|
}
|
|
}
|
|
}
|
|
|
|
private async Task<List<string>?> TryRefreshModelListAsync(
|
|
Func<CancellationToken, Task<List<string>>> refresh,
|
|
string listName,
|
|
CancellationToken cancellationToken)
|
|
{
|
|
try
|
|
{
|
|
using var refreshCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
|
refreshCts.CancelAfter(ModelRefreshTimeout);
|
|
|
|
return await refresh(refreshCts.Token);
|
|
}
|
|
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
throw;
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
_logger.LogWarning("Timed out refreshing {ModelListName} upstream model list.", listName);
|
|
return null;
|
|
}
|
|
catch (Exception exception)
|
|
{
|
|
_logger.LogWarning(
|
|
"Failed to refresh {ModelListName} upstream model list: {Message}",
|
|
listName,
|
|
exception.Message);
|
|
return null;
|
|
}
|
|
}
|
|
|
|
private (List<string> Models, List<string> ActiveModels) UpdateCachedModelSnapshot(
|
|
List<string>? models,
|
|
List<string>? activeModels)
|
|
{
|
|
lock (_modelSnapshotLock)
|
|
{
|
|
if (models is not null)
|
|
{
|
|
_lastModels = models;
|
|
}
|
|
|
|
if (activeModels is not null)
|
|
{
|
|
_lastActiveModels = activeModels;
|
|
}
|
|
|
|
return (
|
|
[.. _lastModels],
|
|
[.. _lastActiveModels]);
|
|
}
|
|
}
|
|
|
|
private static List<string> NormalizeModelNames(IEnumerable<string> models) =>
|
|
models
|
|
.Where(model => !string.IsNullOrWhiteSpace(model))
|
|
.Select(model => model.Trim())
|
|
.Distinct(StringComparer.OrdinalIgnoreCase)
|
|
.OrderBy(model => model, StringComparer.OrdinalIgnoreCase)
|
|
.ToList();
|
|
|
|
private async Task ReceiveLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
|
|
{
|
|
while (socket.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
|
|
{
|
|
var message = await WebSocketMessageTransport.ReceiveAsync(socket, cancellationToken);
|
|
if (message is null)
|
|
{
|
|
break;
|
|
}
|
|
|
|
await DispatchAsync(socket, message, cancellationToken);
|
|
}
|
|
}
|
|
|
|
private Task DispatchAsync(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
|
{
|
|
switch (message.Type)
|
|
{
|
|
case TunnelMessageTypes.HttpRequest:
|
|
_pendingRequestBodies[message.RequestId] = new PendingRequestBody();
|
|
_ = Task.Run(() => StartRequest(socket, message, cancellationToken), cancellationToken);
|
|
break;
|
|
|
|
case TunnelMessageTypes.HttpRequestBody:
|
|
if (_activeRequests.TryGetValue(message.RequestId, out var requestWithBody))
|
|
{
|
|
requestWithBody.AddBody(message.Body ?? []);
|
|
}
|
|
else if (_pendingRequestBodies.TryGetValue(message.RequestId, out var pendingBody))
|
|
{
|
|
pendingBody.AddBody(message.Body ?? []);
|
|
}
|
|
break;
|
|
|
|
case TunnelMessageTypes.HttpRequestComplete:
|
|
if (_activeRequests.TryGetValue(message.RequestId, out var completedRequest))
|
|
{
|
|
completedRequest.CompleteBody();
|
|
}
|
|
else if (_pendingRequestBodies.TryGetValue(message.RequestId, out var pendingBody))
|
|
{
|
|
pendingBody.Complete();
|
|
}
|
|
break;
|
|
|
|
case TunnelMessageTypes.Cancel:
|
|
if (_activeRequests.TryRemove(message.RequestId, out var cancelledRequest))
|
|
{
|
|
cancelledRequest.Cancel();
|
|
}
|
|
else
|
|
{
|
|
_pendingRequestBodies.TryRemove(message.RequestId, out _);
|
|
}
|
|
break;
|
|
|
|
case TunnelMessageTypes.ModelCommand:
|
|
_ = Task.Run(() => RunModelCommandAsync(socket, message, cancellationToken), cancellationToken);
|
|
break;
|
|
}
|
|
|
|
return Task.CompletedTask;
|
|
}
|
|
|
|
private async Task RunModelCommandAsync(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
|
{
|
|
try
|
|
{
|
|
var response = await ExecuteModelCommandAsync(message, cancellationToken);
|
|
if (response.StatusCode is >= 200 and < 300)
|
|
{
|
|
await RefreshModelsOnceAsync(socket, cancellationToken);
|
|
}
|
|
|
|
await SendAsync(socket, response, cancellationToken);
|
|
}
|
|
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
|
|
{
|
|
}
|
|
catch (Exception exception)
|
|
{
|
|
await SendAsync(
|
|
socket,
|
|
new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = message.RequestId,
|
|
Error = exception.Message
|
|
},
|
|
CancellationToken.None);
|
|
}
|
|
}
|
|
|
|
private async Task<TunnelMessage> ExecuteModelCommandAsync(TunnelMessage message, CancellationToken cancellationToken)
|
|
{
|
|
if (string.IsNullOrWhiteSpace(message.RequestId))
|
|
{
|
|
throw new InvalidOperationException("Model command is missing a request id.");
|
|
}
|
|
|
|
if (string.IsNullOrWhiteSpace(message.Model))
|
|
{
|
|
throw new InvalidOperationException("Model command is missing a model name.");
|
|
}
|
|
|
|
if (_llamaCppManager is not null)
|
|
{
|
|
return await ExecuteModelCommandWithLlamaCppAsync(message, cancellationToken);
|
|
}
|
|
|
|
using var request = BuildModelCommandRequest(message);
|
|
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
|
var body = await response.Content.ReadAsByteArrayAsync(cancellationToken);
|
|
|
|
if (ShouldRetryModelCommandWithEmbedding(message.Command, response, body))
|
|
{
|
|
using var embeddingRequest = BuildEmbeddingModelCommandRequest(_options.Upstream, message.Command, message.Model);
|
|
using var embeddingResponse = await _httpClient.SendAsync(embeddingRequest, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
|
var embeddingBody = await embeddingResponse.Content.ReadAsByteArrayAsync(cancellationToken);
|
|
|
|
return BuildModelCommandResult(message.RequestId, embeddingResponse, embeddingBody);
|
|
}
|
|
|
|
return BuildModelCommandResult(message.RequestId, response, body);
|
|
}
|
|
|
|
private async Task<TunnelMessage> ExecuteModelCommandWithLlamaCppAsync(
|
|
TunnelMessage message, CancellationToken cancellationToken)
|
|
{
|
|
var modelName = message.Model?.Trim();
|
|
var normalizedCommand = NormalizeModelCommand(message.Command);
|
|
|
|
switch (normalizedCommand)
|
|
{
|
|
case "load":
|
|
{
|
|
var models = _llamaCppManager!.DiscoverModelsWithBlob();
|
|
var model = models.FirstOrDefault(m =>
|
|
string.Equals(m.OllamaName, modelName, StringComparison.OrdinalIgnoreCase));
|
|
|
|
if (model is null)
|
|
{
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = message.RequestId,
|
|
StatusCode = 404,
|
|
Error = $"Model '{modelName}' not found in Ollama models path."
|
|
};
|
|
}
|
|
|
|
var started = await _llamaCppManager.StartModelContainerAsync(model, cancellationToken);
|
|
if (!started)
|
|
{
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = message.RequestId,
|
|
StatusCode = 500,
|
|
Error = $"Unable to load model '{modelName}'."
|
|
};
|
|
}
|
|
|
|
return BuildModelCommandResult(message.RequestId, 200, "OK", []);
|
|
}
|
|
|
|
case "unload":
|
|
{
|
|
var stopped = await _llamaCppManager!.StopModelContainerAsync(modelName!, cancellationToken);
|
|
if (!stopped)
|
|
{
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = message.RequestId,
|
|
StatusCode = 404,
|
|
Error = $"No running llama.cpp container for model '{modelName}'."
|
|
};
|
|
}
|
|
|
|
return BuildModelCommandResult(message.RequestId, 200, "OK", []);
|
|
}
|
|
|
|
case "pull":
|
|
case "delete":
|
|
{
|
|
using var request = BuildModelCommandRequest(_options.Upstream, message.Command, message.Model);
|
|
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
|
var body = await response.Content.ReadAsByteArrayAsync(cancellationToken);
|
|
return BuildModelCommandResult(message.RequestId, response, body);
|
|
}
|
|
|
|
case "show":
|
|
{
|
|
var models = _llamaCppManager!.DiscoverModelsWithBlob();
|
|
var model = models.FirstOrDefault(m =>
|
|
string.Equals(m.OllamaName, modelName, StringComparison.OrdinalIgnoreCase));
|
|
|
|
if (model is null)
|
|
{
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = message.RequestId,
|
|
StatusCode = 404,
|
|
Error = $"Model '{modelName}' not found in Ollama models path."
|
|
};
|
|
}
|
|
|
|
var showResponse = new
|
|
{
|
|
modelfile = $"# llama.cpp via Docker\nFROM {model.BlobDigest}\n",
|
|
details = new
|
|
{
|
|
format = "gguf",
|
|
family = "llama",
|
|
parameter_size = "",
|
|
quantization_level = ""
|
|
},
|
|
model_info = new { }
|
|
};
|
|
|
|
var body = JsonSerializer.SerializeToUtf8Bytes(showResponse, JsonOptions);
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = message.RequestId,
|
|
StatusCode = 200,
|
|
ReasonPhrase = "OK",
|
|
Body = body
|
|
};
|
|
}
|
|
|
|
default:
|
|
throw new InvalidOperationException($"Unsupported model command '{message.Command}' with llama.cpp.");
|
|
}
|
|
}
|
|
|
|
private static TunnelMessage BuildModelCommandResult(
|
|
string requestId,
|
|
HttpResponseMessage response,
|
|
byte[] body)
|
|
{
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = requestId,
|
|
StatusCode = (int)response.StatusCode,
|
|
ReasonPhrase = response.ReasonPhrase,
|
|
Body = body
|
|
};
|
|
}
|
|
|
|
private static TunnelMessage BuildModelCommandResult(
|
|
string requestId, int statusCode, string reasonPhrase, byte[] body)
|
|
{
|
|
return new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.ModelCommandResult,
|
|
RequestId = requestId,
|
|
StatusCode = statusCode,
|
|
ReasonPhrase = reasonPhrase,
|
|
Body = body
|
|
};
|
|
}
|
|
|
|
private static bool ShouldRetryModelCommandWithEmbedding(
|
|
string? command,
|
|
HttpResponseMessage response,
|
|
byte[] body)
|
|
{
|
|
var normalizedCommand = NormalizeModelCommand(command);
|
|
if (normalizedCommand is not ("load" or "unload")
|
|
|| response.IsSuccessStatusCode)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
if (response.StatusCode == System.Net.HttpStatusCode.BadRequest)
|
|
{
|
|
return true;
|
|
}
|
|
|
|
var responseText = body.Length > 0
|
|
? Encoding.UTF8.GetString(body)
|
|
: "";
|
|
|
|
return responseText.Contains("does not support generate", StringComparison.OrdinalIgnoreCase);
|
|
}
|
|
|
|
private HttpRequestMessage BuildModelCommandRequest(TunnelMessage message)
|
|
{
|
|
return BuildModelCommandRequest(_options.Upstream, message.Command, message.Model);
|
|
}
|
|
|
|
internal static HttpRequestMessage BuildModelCommandRequest(Uri upstream, string? command, string? modelName)
|
|
{
|
|
if (string.IsNullOrWhiteSpace(modelName))
|
|
{
|
|
throw new InvalidOperationException("Model command is missing a model name.");
|
|
}
|
|
|
|
var model = modelName.Trim();
|
|
var normalizedCommand = NormalizeModelCommand(command);
|
|
|
|
return normalizedCommand switch
|
|
{
|
|
"pull" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/pull"))
|
|
{
|
|
Content = JsonContent(new { model, stream = false })
|
|
},
|
|
"delete" => new HttpRequestMessage(HttpMethod.Delete, new Uri(upstream, "/api/delete"))
|
|
{
|
|
Content = JsonContent(new { model })
|
|
},
|
|
"load" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/generate"))
|
|
{
|
|
Content = JsonContent(new { model, stream = false, keep_alive = -1 })
|
|
},
|
|
"unload" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/generate"))
|
|
{
|
|
Content = JsonContent(new { model, stream = false, keep_alive = 0 })
|
|
},
|
|
"show" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/show"))
|
|
{
|
|
Content = JsonContent(new { model })
|
|
},
|
|
_ => throw new InvalidOperationException($"Unsupported model command '{command}'.")
|
|
};
|
|
}
|
|
|
|
internal static HttpRequestMessage BuildEmbeddingModelCommandRequest(Uri upstream, string? command, string? modelName)
|
|
{
|
|
if (string.IsNullOrWhiteSpace(modelName))
|
|
{
|
|
throw new InvalidOperationException("Model command is missing a model name.");
|
|
}
|
|
|
|
var model = modelName.Trim();
|
|
var normalizedCommand = NormalizeModelCommand(command);
|
|
|
|
return normalizedCommand switch
|
|
{
|
|
"load" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/embed"))
|
|
{
|
|
Content = JsonContent(new { model, input = EmbeddingWarmupInput, keep_alive = -1 })
|
|
},
|
|
"unload" => new HttpRequestMessage(HttpMethod.Post, new Uri(upstream, "/api/embed"))
|
|
{
|
|
Content = JsonContent(new { model, input = EmbeddingWarmupInput, keep_alive = 0 })
|
|
},
|
|
_ => throw new InvalidOperationException($"Unsupported embedding model command '{command}'.")
|
|
};
|
|
}
|
|
|
|
private static string NormalizeModelCommand(string? command) =>
|
|
(command ?? "").Trim().ToLowerInvariant();
|
|
|
|
private static StringContent JsonContent<T>(T value) =>
|
|
new(JsonSerializer.Serialize(value, JsonOptions), Encoding.UTF8, "application/json");
|
|
|
|
private async Task StartRequest(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
|
{
|
|
Uri? effectiveUpstream = null;
|
|
var modelName = UpstreamRequest.ExtractModelName(message);
|
|
|
|
if (_llamaCppManager is not null)
|
|
{
|
|
if (modelName is not null)
|
|
{
|
|
effectiveUpstream = _llamaCppManager.GetUpstream(modelName);
|
|
if (effectiveUpstream is not null)
|
|
{
|
|
var running = await _llamaCppManager.IsContainerRunningAsync(modelName);
|
|
if (!running)
|
|
{
|
|
_logger.LogWarning(
|
|
"Cached llama.cpp container for model '{Model}' is not running anymore. Starting a fresh one on demand...",
|
|
modelName);
|
|
effectiveUpstream = null;
|
|
}
|
|
}
|
|
|
|
if (effectiveUpstream is null)
|
|
{
|
|
var model = _llamaCppManager.DiscoverModelsWithBlob()
|
|
.FirstOrDefault(m => string.Equals(m.OllamaName, modelName, StringComparison.OrdinalIgnoreCase));
|
|
|
|
if (model is not null)
|
|
{
|
|
_logger.LogInformation(
|
|
"Request for model '{Model}' but no llama.cpp container is running. Starting one on demand...",
|
|
modelName);
|
|
|
|
var started = await _llamaCppManager.StartModelContainerAsync(model, cancellationToken);
|
|
if (started)
|
|
{
|
|
effectiveUpstream = _llamaCppManager.GetUpstream(modelName);
|
|
}
|
|
else
|
|
{
|
|
await SendModelLoadErrorAsync(socket, message, modelName, cancellationToken);
|
|
return;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
_logger.LogWarning(
|
|
"Model '{Model}' was not found in the Ollama models path. Falling back to default upstream.",
|
|
modelName);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// When llama.cpp backend is active, set up request/response translation
|
|
Func<string, string?>? pathTransform = null;
|
|
Func<byte[], byte[]>? bodyTransform = null;
|
|
Func<HttpResponseMessage, CancellationToken, Task>? responseHandler = null;
|
|
string? translatorModelName = null;
|
|
|
|
if (_llamaCppManager is not null && effectiveUpstream is not null && modelName is not null)
|
|
{
|
|
var translator = new OllamaToLlamaCppTranslator(modelName, _logger);
|
|
translatorModelName = modelName;
|
|
bool originalRequestedStream = true;
|
|
|
|
pathTransform = path =>
|
|
{
|
|
if (translator.TryTranslatePath(message.Method ?? "GET", path, out var newPath))
|
|
{
|
|
return newPath;
|
|
}
|
|
return null;
|
|
};
|
|
|
|
bodyTransform = body =>
|
|
{
|
|
originalRequestedStream = OllamaToLlamaCppTranslator.ExtractOriginalStream(body, message.PathAndQuery ?? "/");
|
|
return translator.TranslateBody(message.PathAndQuery ?? "/", body);
|
|
};
|
|
|
|
responseHandler = translator.CreateResponseHandler(
|
|
(response, token) => SendAsync(socket, response, token),
|
|
message.RequestId,
|
|
message.PathAndQuery ?? "/",
|
|
() => originalRequestedStream);
|
|
}
|
|
|
|
var request = new UpstreamRequest(
|
|
_options,
|
|
_httpClient,
|
|
message,
|
|
(response, token) => SendAsync(socket, response, token),
|
|
requestId => _activeRequests.TryRemove(requestId, out _),
|
|
cancellationToken,
|
|
effectiveUpstream: effectiveUpstream,
|
|
responseHandler: responseHandler,
|
|
pathTransform: pathTransform,
|
|
bodyTransform: bodyTransform,
|
|
onConnectionRefused: () =>
|
|
{
|
|
if (_llamaCppManager is not null && modelName is not null)
|
|
{
|
|
_llamaCppManager.RemoveModelMapping(modelName);
|
|
}
|
|
});
|
|
|
|
if (!_activeRequests.TryAdd(message.RequestId, request))
|
|
{
|
|
_ = SendAsync(
|
|
socket,
|
|
new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.Error,
|
|
RequestId = message.RequestId,
|
|
Error = "Duplicate request id."
|
|
},
|
|
cancellationToken);
|
|
_pendingRequestBodies.TryRemove(message.RequestId, out _);
|
|
return;
|
|
}
|
|
|
|
if (_pendingRequestBodies.TryRemove(message.RequestId, out var pendingBody))
|
|
{
|
|
pendingBody.TransferTo(request);
|
|
}
|
|
|
|
await request.RunAsync();
|
|
}
|
|
|
|
private async Task SendAsync(ClientWebSocket socket, TunnelMessage message, CancellationToken cancellationToken)
|
|
{
|
|
await _sendLock.WaitAsync(cancellationToken);
|
|
try
|
|
{
|
|
if (socket.State == WebSocketState.Open)
|
|
{
|
|
await WebSocketMessageTransport.SendAsync(socket, message, cancellationToken);
|
|
}
|
|
}
|
|
finally
|
|
{
|
|
_sendLock.Release();
|
|
}
|
|
}
|
|
|
|
private async Task SendModelLoadErrorAsync(
|
|
ClientWebSocket socket, TunnelMessage message, string modelName, CancellationToken cancellationToken)
|
|
{
|
|
_logger.LogWarning(
|
|
"Unable to load model '{Model}' via llama.cpp. Notifying caller.", modelName);
|
|
|
|
var body = JsonSerializer.SerializeToUtf8Bytes(
|
|
new { error = $"Unable to load model '{modelName}'" },
|
|
JsonOptions);
|
|
|
|
_pendingRequestBodies.TryRemove(message.RequestId, out _);
|
|
|
|
await SendAsync(
|
|
socket,
|
|
new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.HttpResponseHeaders,
|
|
RequestId = message.RequestId,
|
|
StatusCode = 500,
|
|
ReasonPhrase = "Internal Server Error",
|
|
Headers = [new HeaderPair("Content-Type", "application/json")]
|
|
},
|
|
CancellationToken.None);
|
|
|
|
await SendAsync(
|
|
socket,
|
|
new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.HttpResponseBody,
|
|
RequestId = message.RequestId,
|
|
Body = body
|
|
},
|
|
CancellationToken.None);
|
|
|
|
await SendAsync(
|
|
socket,
|
|
new TunnelMessage
|
|
{
|
|
Type = TunnelMessageTypes.HttpResponseComplete,
|
|
RequestId = message.RequestId
|
|
},
|
|
CancellationToken.None);
|
|
}
|
|
|
|
private void CancelAllActiveRequests()
|
|
{
|
|
foreach (var pair in _activeRequests.ToArray())
|
|
{
|
|
if (_activeRequests.TryRemove(pair.Key, out var request))
|
|
{
|
|
request.Cancel();
|
|
}
|
|
}
|
|
|
|
_pendingRequestBodies.Clear();
|
|
}
|
|
}
|
|
|
|
internal sealed class PendingRequestBody
|
|
{
|
|
private readonly List<byte[]> _chunks = [];
|
|
private bool _completed;
|
|
private readonly object _lock = new();
|
|
|
|
public void AddBody(byte[] chunk)
|
|
{
|
|
if (chunk.Length == 0)
|
|
return;
|
|
|
|
lock (_lock)
|
|
{
|
|
_chunks.Add(chunk);
|
|
}
|
|
}
|
|
|
|
public void Complete()
|
|
{
|
|
lock (_lock)
|
|
{
|
|
_completed = true;
|
|
}
|
|
}
|
|
|
|
public void TransferTo(UpstreamRequest request)
|
|
{
|
|
lock (_lock)
|
|
{
|
|
foreach (var chunk in _chunks)
|
|
{
|
|
request.AddBody(chunk);
|
|
}
|
|
|
|
_chunks.Clear();
|
|
|
|
if (_completed)
|
|
{
|
|
request.CompleteBody();
|
|
}
|
|
}
|
|
}
|
|
}
|