MemoryEmbeddings.cs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

using System.Net.Http.Json;
using System.Text.Json;

namespace Lilith.Core.Memory;

public interface IMemoryEmbeddings
{
    Task<string> IdentityAsync(CancellationToken ct);
    Task<float[][]> EmbedAsync(string[] inputs, CancellationToken ct);
}

internal sealed class OllamaMemoryEmbeddings : IMemoryEmbeddings, IDisposable
{
    private const string Model = "embeddinggemma:latest";
    private readonly HttpClient http = new(new HttpClientHandler { UseProxy = false, AllowAutoRedirect = false })
        { BaseAddress = new Uri("http://127.0.0.1:11434/"), Timeout = TimeSpan.FromSeconds(60), MaxResponseContentBufferSize = 8_000_000 };

    public async Task<string> IdentityAsync(CancellationToken ct)
    {
        using var response = await http.GetAsync("api/tags", ct);
        response.EnsureSuccessStatusCode();
        using var json = JsonDocument.Parse(await response.Content.ReadAsStringAsync(ct));
        foreach (var model in json.RootElement.GetProperty("models").EnumerateArray())
            if (model.GetProperty("name").GetString() == Model)
            {
                string digest = model.GetProperty("digest").GetString() ?? "";
                if (!System.Text.RegularExpressions.Regex.IsMatch(digest, "^[A-Fa-f0-9]{64}$")) throw new InvalidDataException("Embedding model identity is invalid.");
                return Model + ":" + digest;
            }
        throw new IOException("EmbeddingGemma is not installed. Install it explicitly with ollama pull embeddinggemma.");
    }

    public async Task<float[][]> EmbedAsync(string[] inputs, CancellationToken ct)
    {
        using var response = await http.PostAsJsonAsync("api/embed", new { model = Model, input = inputs, truncate = false, keep_alive = "0s", options = new { num_ctx = 2048 } }, ct);
        response.EnsureSuccessStatusCode();
        using var json = JsonDocument.Parse(await response.Content.ReadAsStringAsync(ct));
        return json.RootElement.GetProperty("embeddings").EnumerateArray()
            .Select(row => row.EnumerateArray().Select(v => v.GetSingle()).ToArray()).ToArray();
    }
    public void Dispose() => http.Dispose();
}