TentacleBot/Services/LlmService.cs

using System.Collections.Concurrent;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Configuration;

namespace TentacleBot.Services;

public class LlmService
{
    private readonly HttpClient _httpClient;
    private readonly ILogger<LlmService> _logger;
    private readonly string _endpoint;
    private readonly string _apiKey;
    private readonly string _model;
    private readonly int _maxTokens;
    private readonly double _temperature;
    private readonly ConcurrentDictionary<ulong, List<ChatMessage>> _histories = new();
    private const int MaxHistoryChars = 20000;

    public LlmService(IConfiguration configuration, ILogger<LlmService> logger)
    {
        _logger = logger;
        _httpClient = new HttpClient();

        _endpoint = configuration["Llm:Endpoint"] ?? throw new InvalidOperationException("Llm:Endpoint is not configured");
        _apiKey = Environment.GetEnvironmentVariable("LLM_API_KEY")
            ?? throw new InvalidOperationException("LLM_API_KEY environment variable is not set");
        _model = configuration["Llm:Model"] ?? throw new InvalidOperationException("Llm:Model is not configured");
        _maxTokens = int.Parse(configuration["Llm:MaxTokens"] ?? "1024");
        _temperature = double.Parse(configuration["Llm:Temperature"] ?? "0.7");

        _httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", _apiKey);
        _httpClient.DefaultRequestHeaders.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
    }

    public async IAsyncEnumerable<string> SendMessageAsync(ulong channelId, string user, string message)
    {
        var history = _histories.GetOrAdd(channelId, _ => new List<ChatMessage>());

        history.Add(new ChatMessage { Role = "user", Content = message });

        TrimHistory(history);

        var request = new ChatRequest
        {
            Model = _model,
            Messages = history.Select(m => new ChatMessage { Role = m.Role, Content = m.Content }).ToList(),
            MaxTokens = _maxTokens,
            Temperature = (float)_temperature,
            Stream = true
        };

        var json = JsonSerializer.Serialize(request, new JsonSerializerOptions { PropertyNamingPolicy = JsonNamingPolicy.CamelCase });
        var content = new StringContent(json, Encoding.UTF8, "application/json");

        using var response = await _httpClient.PostAsync(_endpoint, content);

        if (!response.IsSuccessStatusCode)
        {
            var errorBody = await response.Content.ReadAsStringAsync();
            _logger.LogError("LLM API error ({StatusCode}): {Error}", response.StatusCode, errorBody);
            throw new HttpRequestException($"LLM API returned {response.StatusCode}: {errorBody}");
        }

        var accumulated = new StringBuilder();

        var mediaType = response.Content.Headers.ContentType?.MediaType;
        if (string.Equals(mediaType, "text/event-stream", StringComparison.OrdinalIgnoreCase))
        {
            await using var stream = await response.Content.ReadAsStreamAsync();
            using var reader = new StreamReader(stream);

            while (await reader.ReadLineAsync() is { } line)
            {
                if (!line.StartsWith("data:", StringComparison.Ordinal))
                    continue;

                var data = line.Substring(5).Trim();
                if (data == "[DONE]")
                    break;

                var token = TryParseStreamToken(data);
                if (!string.IsNullOrEmpty(token))
                {
                    accumulated.Append(token);
                    yield return accumulated.ToString();
                }
            }
        }
        else
        {
            var body = await response.Content.ReadAsStringAsync();
            var fullText = TryParseNonStreamContent(body);
            if (!string.IsNullOrWhiteSpace(fullText))
            {
                accumulated.Append(fullText);
                yield return accumulated.ToString();
            }
        }

        if (accumulated.Length > 0)
        {
            history.Add(new ChatMessage { Role = "assistant", Content = accumulated.ToString() });
            TrimHistory(history);
        }
    }

    private static string? TryParseStreamToken(string data)
    {
        try
        {
            using var doc = JsonDocument.Parse(data);

            if (doc.RootElement.TryGetProperty("choices", out var choices) &&
                choices.ValueKind == JsonValueKind.Array &&
                choices.GetArrayLength() > 0)
            {
                var firstChoice = choices[0];

                if (firstChoice.TryGetProperty("delta", out var delta) &&
                    delta.ValueKind == JsonValueKind.Object &&
                    delta.TryGetProperty("content", out var deltaContent) &&
                    deltaContent.ValueKind == JsonValueKind.String)
                {
                    return deltaContent.GetString();
                }

                if (firstChoice.TryGetProperty("text", out var text) &&
                    text.ValueKind == JsonValueKind.String)
                {
                    return text.GetString();
                }
            }

            if (doc.RootElement.TryGetProperty("delta", out var rootDelta) &&
                rootDelta.ValueKind == JsonValueKind.Object &&
                rootDelta.TryGetProperty("content", out var rootContent) &&
                rootContent.ValueKind == JsonValueKind.String)
            {
                return rootContent.GetString();
            }

            if (doc.RootElement.TryGetProperty("content", out var content) &&
                content.ValueKind == JsonValueKind.String)
            {
                return content.GetString();
            }

            return null;
        }
        catch
        {
            return null;
        }
    }

    private static string? TryParseNonStreamContent(string body)
    {
        try
        {
            using var doc = JsonDocument.Parse(body);

            if (doc.RootElement.TryGetProperty("choices", out var choices) &&
                choices.ValueKind == JsonValueKind.Array &&
                choices.GetArrayLength() > 0)
            {
                var firstChoice = choices[0];

                if (firstChoice.TryGetProperty("message", out var message) &&
                    message.ValueKind == JsonValueKind.Object &&
                    message.TryGetProperty("content", out var messageContent) &&
                    messageContent.ValueKind == JsonValueKind.String)
                {
                    return messageContent.GetString();
                }

                if (firstChoice.TryGetProperty("text", out var text) &&
                    text.ValueKind == JsonValueKind.String)
                {
                    return text.GetString();
                }
            }

            if (doc.RootElement.TryGetProperty("content", out var content) &&
                content.ValueKind == JsonValueKind.String)
            {
                return content.GetString();
            }

            return null;
        }
        catch
        {
            return null;
        }
    }

    public void ResetHistory(ulong channelId)
    {
        _histories.TryRemove(channelId, out _);
    }

    private void TrimHistory(List<ChatMessage> history)
    {
        while (history.Sum(m => m.Content.Length) > MaxHistoryChars && history.Count > 2)
        {
            history.RemoveAt(0);
        }
    }

    private record ChatRequest
    {
        public string Model { get; init; } = "";
        public List<ChatMessage> Messages { get; init; } = new();
        public int MaxTokens { get; init; }
        public float Temperature { get; init; }
        public bool Stream { get; init; }
    }

    public record ChatMessage
    {
        public string Role { get; set; } = "";
        public string Content { get; set; } = "";
    }
}