Luna/Luna.Providers/Mistral/MistralChatClient.cs

using System.Runtime.CompilerServices;
using System.Text.Json;
using Luna.Providers.Mistral.DTOs;
using Microsoft.Extensions.AI;

namespace Luna.Providers.Mistral;

/// <summary>
/// Custom IChatClient implementation for the Mistral API.
/// Implements chat completion according to the Mistral OpenAPI spec.
/// </summary>
public class MistralChatClient : IChatClient
{
    private readonly HttpClient httpClient;
    private readonly string model;
    private bool disposed;

    public MistralChatClient(HttpClient httpClient, string model)
    {
        this.httpClient = httpClient ?? throw new ArgumentNullException(nameof(httpClient));
        this.model = model ?? throw new ArgumentNullException(nameof(model));
    }

    public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
    {
        var request = CreateChatCompletionRequest(messages, options);
        var response = await SendChatCompletionRequest(request, cancellationToken);
        return MapToChatResponse(response);
    }

    public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        var request = CreateChatCompletionRequest(messages, options, true);

        await foreach (var update in SendStreamingRequest(request, cancellationToken))
        {
            yield return MapToChatResponseUpdate(update);
        }
    }

    public object? GetService(Type serviceType, object? serviceKey = null)
    {
        return serviceType == typeof(HttpClient) 
            ? httpClient 
            : null;
    }

    public ChatClientMetadata Metadata => new(
        "Mistral",
        httpClient.BaseAddress,
        model
    );

    public void Dispose()
    {
        GC.SuppressFinalize(this);
        if (disposed) return;
        
        httpClient.Dispose();
        disposed = true;
    }

    /// <summary>
    /// Creates a chat completion request according to the Mistral API spec.
    /// Reference: /v1/chat/completions endpoint in the OpenAPI spec.
    /// </summary>
    private MistralChatCompletionRequest CreateChatCompletionRequest(IEnumerable<ChatMessage> messages, ChatOptions? options, bool streaming = false)
    {
        var messagesWithInstructions = new List<ChatMessage>();
        if (!string.IsNullOrWhiteSpace(options?.Instructions))
        {
            messagesWithInstructions.Add(new ChatMessage(ChatRole.System, options.Instructions));
        }
        messagesWithInstructions.AddRange(messages);
        
        var request = new MistralChatCompletionRequest
        {
            Model = options?.ModelId ?? model,
            Messages = messagesWithInstructions.Select(m => new MistralMessage
            {
                Role = m.Role.ToString().ToLowerInvariant(),
                Content = m.Text
            }).ToList(),
            Temperature = options?.Temperature,
            MaxTokens = options?.MaxOutputTokens,
            TopP = options?.TopP,
            Stream = streaming,
            ResponseFormat = options?.ResponseFormat is ChatResponseFormatJson
                ? new ResponseFormat { Type = "json_object" }
                : null
        };

        return request;
    }

    private async Task<MistralChatCompletionResponse> SendChatCompletionRequest(
        MistralChatCompletionRequest request,
        CancellationToken cancellationToken)
    {
        var json = JsonSerializer.Serialize(request, MistralJsonContext.Default.MistralChatCompletionRequest);
        var content = new StringContent(json, System.Text.Encoding.UTF8, "application/json");

        var response = await httpClient.PostAsync("chat/completions", content, cancellationToken);
        response.EnsureSuccessStatusCode();

        var responseJson = await response.Content.ReadAsStringAsync(cancellationToken);
        return JsonSerializer.Deserialize(responseJson, MistralJsonContext.Default.MistralChatCompletionResponse)
            ?? throw new InvalidOperationException("Failed to deserialize response");
    }

    private async IAsyncEnumerable<MistralChatCompletionChunk> SendStreamingRequest(
        MistralChatCompletionRequest request,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var json = JsonSerializer.Serialize(request, MistralJsonContext.Default.MistralChatCompletionRequest);
        var content = new StringContent(json, System.Text.Encoding.UTF8, "application/json");

        var response = await httpClient.PostAsync("chat/completions", content, cancellationToken);
        response.EnsureSuccessStatusCode();

        var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
        using var reader = new StreamReader(stream);

        while (!cancellationToken.IsCancellationRequested)
        {
            var line = await reader.ReadLineAsync(cancellationToken);
            if (string.IsNullOrWhiteSpace(line) || !line.StartsWith("data: "))
                continue;

            var data = line[6..]; // Remove "data: " prefix
            if (data == "[DONE]")
                yield break;

            var chunk = JsonSerializer.Deserialize(data, MistralJsonContext.Default.MistralChatCompletionChunk);
            if (chunk != null)
            {
                yield return chunk;
            }
        }
    }

    private static ChatResponse MapToChatResponse(MistralChatCompletionResponse response)
    {
        var choice = response.Choices?.FirstOrDefault();
        if (choice?.Message == null)
        {
            throw new InvalidOperationException("No choices in response");
        }

        return new ChatResponse(
            new ChatMessage(
                MapRole(choice.Message.Role),
                choice.Message.Content
            )
        );
    }

    private static ChatResponseUpdate MapToChatResponseUpdate(MistralChatCompletionChunk chunk)
    {
        var choice = chunk.Choices?.FirstOrDefault();
        if (choice?.Delta == null)
        {
            return new ChatResponseUpdate();
        }

        return new ChatResponseUpdate(
            MapRole(choice.Delta.Role),
            choice.Delta.Content
        );
    }

    private static ChatRole MapRole(string? role) => role?.ToLowerInvariant() switch
    {
        "system" => ChatRole.System,
        "user" => ChatRole.User,
        "assistant" => ChatRole.Assistant,
        "tool" => ChatRole.Tool,
        _ => ChatRole.User
    };
}