LMStudioClient.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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

using System.Net.Http.Headers;
using System.Text.Json;
using Lilith.Core.Tools;

namespace Lilith.Core.LMStudio;

public sealed class LMStudioClient(Logger logger, ChatHistoryStore? historyStore = null, HttpClient? httpClient = null)
    : ProviderClient(logger, historyStore, httpClient)
{
    public string BaseUrl { get; init; } = LMStudioSetup.BaseUrl;

    protected override async Task<ChatChunk?> GenerateReplyAsync(List<ChatMessage> messages)
        => await GenerateReplyAsync(messages, CancellationToken.None);

    protected override async Task<ChatChunk?> GenerateReplyAsync(List<ChatMessage> messages, CancellationToken cancellationToken)
    {
        if (messages.Any(m => m.Images?.Count > 0))
            throw new NotSupportedException("LM Studio image input is not implemented in this interface.");
        using var request = HttpExtensions.JsonRequest(BaseUrl + "/v1/chat/completions", new
        {
            model = Model,
            stream = false,
            messages = messages.Select(ToWireMessage),
            tools = ToolDefinitions
        });
        LMStudioSetup.Authorize(request);
        using var json = await Http.SendJsonDocumentAsync(request, cancellationToken);
        var choice = json.RootElement.GetProperty("choices")[0];
        string? finish = choice.GetProperty("finish_reason").GetString();
        if (finish is not ("stop" or "tool_calls"))
            throw new InvalidDataException("LM Studio did not return a complete text reply.");
        var message = choice.GetProperty("message");
        if (message.GetProperty("role").GetString() != "assistant")
            throw new InvalidDataException("LM Studio returned an invalid assistant message.");
        List<ToolCall>? calls = null;
        if (message.TryGetProperty("tool_calls", out var toolCalls) && toolCalls.ValueKind != JsonValueKind.Null)
        {
            calls = [];
            foreach (var call in toolCalls.EnumerateArray())
            {
                if (call.GetProperty("type").GetString() != "function" ||
                    string.IsNullOrWhiteSpace(call.GetProperty("id").GetString()))
                    throw new InvalidDataException("Invalid LM Studio function call.");
                var function = call.GetProperty("function");
                using var arguments = JsonDocument.Parse(function.GetProperty("arguments").GetString()!);
                calls.Add(new ToolCall(new ToolFunction(function.GetProperty("name").GetString()!, arguments.RootElement.Clone()),
                    call.GetProperty("id").GetString()));
            }
            if (calls.Select(c => c.Id).Distinct().Count() != calls.Count)
                throw new InvalidDataException("Duplicate LM Studio tool call IDs.");
        }
        if ((finish == "tool_calls") != (calls is { Count: > 0 }))
            throw new InvalidDataException("LM Studio tool calls do not match the finish reason.");
        string? content = message.TryGetProperty("content", out var text) && text.ValueKind == JsonValueKind.String ? text.GetString() : null;
        if (calls is not { Count: > 0 } && content is null)
            throw new InvalidDataException("LM Studio returned no assistant content.");
        int? prompt = null, completion = null;
        if (json.RootElement.TryGetProperty("usage", out var usage))
        {
            if (usage.TryGetProperty("prompt_tokens", out var p)) prompt = p.GetInt32();
            if (usage.TryGetProperty("completion_tokens", out var c)) completion = c.GetInt32();
        }
        return new ChatChunk(new ChatMessage("assistant", content ?? string.Empty, ToolCalls: calls), true, prompt, completion);
    }

    private static object ToWireMessage(ChatMessage message)
    {
        var wire = new Dictionary<string, object?> { ["role"] = message.Role, ["content"] = message.Content };
        if (message.ToolCallId is { } id) wire["tool_call_id"] = id;
        if (message.ToolCalls is { Count: > 0 } calls)
            wire["tool_calls"] = calls.Select(c => new
            {
                id = c.Id, type = "function",
                function = new { name = c.Function.Name, arguments = c.Function.Arguments.GetRawText() }
            }).ToArray();
        return wire;
    }

    public override ProviderClient CreateScenarioBranch(string prompt, Logger branchLogger, ChatHistoryStore store)
    {
        var child = new LMStudioClient(branchLogger, store, Http)
        {
            Model = Model, BaseUrl = BaseUrl, SystemPrompt = prompt, NumCtx = NumCtx, AgentName = "Scenario"
        };
        child.InitializeConversation();
        return child;
    }

    public static async Task<LMStudioClient> CreateClient(string model, string systemPrompt, Logger logger,
        bool startServerIfNotRunning, string agentName = "Lilith", ChatHistoryStore? historyStore = null)
    {
        await LMStudioSetup.EnsureServerAsync(logger, startServerIfNotRunning);
        var client = new LMStudioClient(logger, historyStore) { Model = model, SystemPrompt = systemPrompt, AgentName = agentName };
        try
        {
            var models = await LMStudioSetup.ListModelsAsync();
            var selected = models.FirstOrDefault(m => m.Id == model)
                ?? throw new InvalidOperationException($"LM Studio model '{model}' is not downloaded. Run setup first.");
            if (selected.InstanceId is not null && selected.ContextLength is > 0)
            {
                client.Model = selected.InstanceId;
                client.NumCtx = selected.ContextLength.Value;
            }
            else
            {
                using var request = HttpExtensions.JsonRequest(client.BaseUrl + "/api/v1/models/load", new { model, echo_load_config = true });
                LMStudioSetup.Authorize(request);
                using var json = await client.Http.SendJsonDocumentAsync(request);
                client.Model = json.RootElement.GetProperty("instance_id").GetString()!;
                client.NumCtx = json.RootElement.GetProperty("load_config").GetProperty("context_length").GetInt32();
            }
            if (client.NumCtx <= 0) throw new InvalidDataException("LM Studio reported an invalid context window.");
            client.InitializeConversation();
            return client;
        }
        catch { client.Dispose(); throw; }
    }
}