Bring back AsIChatClient for OpenAI AssistantClient

[stephentoub]

cc: @markwallace-microsoft

Microsoft Reviewers: Open in CodeFlow

[asinghca]

I just wanted to thank you all so much for this update!

I have recently discovered the two AI packages :

https://github.com/openai/openai-dotnet/blob/main/src/Custom/Assistants/AssistantClient.cs
https://github.com/dotnet/extensions/blob/main/src/Libraries/Microsoft.Extensions.AI.OpenAI/OpenAIAssistantChatClient.cs

Leveraging your updated chat client for assistant and the Open AI's assistant API client, I was able to replace a whole bunch of code with just this wrapper to meet my app level interface ILargeLanguageModelService. Prior to this, I had written my own implementation for all of it.

I am curious to know if I am using, the OpenAIAssistantChatClient : IChatClient correctly? I would appreciate any feedback as I have not able to find much documentation on this particular chat client other than the source code for it.

using UseCases.Models.Chat;

namespace UseCases.Interfaces.Chat;

/// <summary>
/// Defines a vendor-neutral contract for interacting with large language model (LLM) providers
/// using application-managed run identifiers. Abstracts the full lifecycle of response generation.
/// </summary>
public interface ILargeLanguageModelService
{
    /// <summary>
    /// Initiates a background response generation task for the specified chat thread and run ID.
    /// Implementations may use <see cref="IChatMessageStore"/> internally to read message context.
    /// </summary>
    /// <param name="chatId">The ID of the chat thread.</param>
    /// <param name="runId">The application-generated identifier for this response generation task.</param>
    /// <param name="ct">Optional cancellation token.</param>
    Task SendMessageAsync(string chatId, string runId, CancellationToken ct = default);

    /// <summary>
    /// Returns the current lifecycle status of the generation task.
    /// </summary>
    /// <param name="chatId">The chat thread ID.</param>
    /// <param name="runId">The application-assigned run ID.</param>
    /// <param name="ct">Optional cancellation token.</param>
    /// <returns>The current status of the generation process.</returns>
    Task<ResponseGenerationStatus> GetStatusAsync(string chatId, string runId, CancellationToken ct = default);

    /// <summary>
    /// Returns the most recent response content generated so far for the given task.
    /// If the task is still InProgress, this may return partial content or an empty string.
    /// If the task is Completed, the content is expected to be final.
    /// Returns null if no content has been generated yet or if the task failed.
    /// </summary>
    /// <param name="chatId">The chat thread ID.</param>
    /// <param name="runId">The application-assigned run ID.</param>
    /// <param name="ct">Optional cancellation token.</param>
    /// <returns>Partial or complete response content, or null if none is available.</returns>
    Task<string?> GetResponseAsync(string chatId, string runId, CancellationToken ct = default);

    /// <summary>
    /// Cancels a background response generation task if it is still in progress.
    /// </summary>
    /// <param name="chatId">The chat thread ID.</param>
    /// <param name="runId">The application-assigned run ID.</param>
    /// <param name="ct">Optional cancellation token.</param>
    Task CancelAsync(string chatId, string runId, CancellationToken ct = default);
}

using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;

using Infrastructure.ApiClient.OpenAI.Assistant;
using UseCases.Interfaces.Auth;
using UseCases.Interfaces.Chat;
using UseCases.Models.Chat;

using AIChatMessage = Microsoft.Extensions.AI.ChatMessage;
using ChatRole = Microsoft.Extensions.AI.ChatRole;
using Microsoft.Extensions.Logging;

namespace Infrastructure.Services.LargeLanguageModel;

/// <summary>
/// Modern implementation of <see cref="ILargeLanguageModelService"/> using Microsoft.Extensions.AI abstractions.
/// Leverages <see cref="IChatClient"/> for OpenAI Assistant interactions with built-in thread management.
/// </summary>
public class OpenAIAssistantService : ILargeLanguageModelService
{
    private readonly IChatClient _chatClient;
    private readonly IAssistantThreadResponseStore _responseStore;
    private readonly IUserContext _user; // We might need this later
    private readonly IChatMessageStore _messageStore;
    private readonly ILogger<OpenAIAssistantService> _logger;
    /// <summary>
    /// Initializes a new instance of the <see cref="OpenAIAssistantService"/> class.
    /// </summary>
    /// <param name="serviceProvider">Service provider to resolve keyed IChatClient.</param>
    /// <param name="responseStore">Backing store for tracking OpenAI chat response state.</param>
    /// <param name="user">Context for the currently authenticated user.</param>
    /// <param name="messageStore">Chat message repository scoped by user and chat ID.</param>
    public OpenAIAssistantService(
        IChatClient chatClient,
        IAssistantThreadResponseStore responseStore,
        IUserContext user,
        IChatMessageStore messageStore,
        ILogger<OpenAIAssistantService> logger)
    {
        _chatClient = chatClient;
        _responseStore = responseStore;
        _user = user;
        _messageStore = messageStore;
        _logger = logger;
    }

    /// <inheritdoc />
    /// <remarks>
    /// 1. Only the latest user message is sent to the AI client.
    /// 2. It's a turn based conversation, so the AI client will respond to the latest user's message.
    /// 3. Since this is an Open AI Assistant client, the AI client is stateful and will remember the conversation.
    /// 4. There is no need to send the entire conversation history to the AI client.
    /// 5. Since we are also wrapping the <see cref="IChatClient"/> with a <see cref="FunctionInvokingChatClient"/>,
    /// the function calling is handled by the <see cref="FunctionInvokingChatClient"/>, so we don't need to handle it ourself.
    /// <see cref="ChatResponse.ConversationId"/> is used to track the created conversation state.
    /// <see cref="ChatOptions.ConversationId"/> is used to track the existing conversation state.
    /// </remarks>
    public async Task SendMessageAsync(string appChatId, string appRunId, CancellationToken ct = default)
    {
        var latest = await Task.Run(() =>
            _messageStore.GetLatestUserMessageForChat(appChatId), ct);
            

        if (latest == null)
            throw new InvalidOperationException("No user message found to respond to.");

        // Convert chat history to Microsoft.Extensions.AI format
        AIChatMessage? chatMessages = ConvertToChatMessages(latest);
        string? openAiThreadId = null;
        _responseStore.TryGetThreadId(appChatId, out openAiThreadId);

        // Ensure empty strings are treated as null to allow OpenAI to create new threads
        if (string.IsNullOrWhiteSpace(openAiThreadId))
        {
            openAiThreadId = null;
        }

        ChatOptions chatOptions = new ChatOptions
        {
            ConversationId = openAiThreadId
        };

        _logger.LogInformation("Starting response generation for chat {AppChatId} and run {AppRunId} with thread ID {ThreadId}", 
            appChatId, appRunId, openAiThreadId ?? "new thread");

        // Create cancellation token source for this specific task
        var cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(ct);
        
        try
        {
            // Create and store the task execution state
            var executionState = new AiGenerationTaskState
            {
                Status = ResponseGenerationStatus.InProgress,
                CancellationTokenSource = cancellationTokenSource,
                StartedAtUtc = DateTime.UtcNow
            };

            // Start the assistant conversation asynchronously
            var backgroundTask = Task.Factory.StartNew(async () =>
            {
                try
                {
                    ChatResponse response = await _chatClient.GetResponseAsync(chatMessages, chatOptions, cancellationToken: cancellationTokenSource.Token);
                    
                    _responseStore.SaveChatResponse(appChatId, appRunId, response);

                    // Save thread ID if present in response (indicates new thread was created or existing thread was used)
                    if(response.ConversationId != null)
                    {
                        _responseStore.SaveThreadId(appChatId, response.ConversationId);
                        _logger.LogInformation("Saved thread ID {ThreadId} for chat {AppChatId}", response.ConversationId, appChatId);
                    }
                    
                    _responseStore.UpdateTaskStatus(appChatId, appRunId, ResponseGenerationStatus.Completed);
                }
                catch (OperationCanceledException)
                {
                    _logger.LogWarning("Response generation was cancelled for chat {AppChatId} and run {AppRunId}", appChatId, appRunId);
                    _responseStore.UpdateTaskStatus(appChatId, appRunId, ResponseGenerationStatus.Cancelled, "Operation was cancelled");
                }
                catch (Exception ex)
                {
                    _logger.LogError("Response generation failed for chat {AppChatId} and run {AppRunId}: {ErrorMessage}", appChatId, appRunId, ex.Message);
                    _responseStore.UpdateTaskStatus(appChatId, appRunId, ResponseGenerationStatus.Failed, ex.Message);
                }
            },
            cancellationTokenSource.Token,
            TaskCreationOptions.LongRunning | TaskCreationOptions.DenyChildAttach,
            TaskScheduler.Default);

            // Store the task execution state with the actual task
            executionState.ExecutingTask = backgroundTask;
            _responseStore.SaveTaskExecutionState(appChatId, appRunId, executionState);
        }
        catch (Exception ex)
        {
            _logger.LogError("Failed to start response generation for chat {AppChatId} and run {AppRunId}: {ErrorMessage}", appChatId, appRunId, ex.Message);
            _responseStore.UpdateTaskStatus(appChatId, appRunId, ResponseGenerationStatus.Failed, "Failed to start generation task");
            cancellationTokenSource.Dispose();
            throw;
        }
    }

    /// <inheritdoc />
    public async Task<ResponseGenerationStatus> GetStatusAsync(string appChatId, string appRunId, CancellationToken ct = default)
    {
        if (!_responseStore.TryGetTaskExecutionState(appChatId, appRunId, out AiGenerationTaskState? executionState))
            return ResponseGenerationStatus.Unknown;

        return await Task.FromResult(executionState!.Status);
    }

    /// <inheritdoc />
    public async Task<string?> GetResponseAsync(string appChatId, string appRunId, CancellationToken ct = default)
    {
        string? responseContent = null;

        if (_responseStore.TryGetChatResponse(appChatId, appRunId, out ChatResponse? response) 
        && response != null
        && response.Messages != null
        && response.Messages.Count > 0) {
            responseContent = response.Messages.Last().ToString();
        }

        return await Task.FromResult(responseContent);
    }

    /// <inheritdoc />
    public async Task CancelAsync(string appChatId, string appRunId, CancellationToken ct = default)
    {
        if (!_responseStore.TryGetTaskExecutionState(appChatId, appRunId, out AiGenerationTaskState? executionState))
            return;

        // Cancel the task if it's still running
        if (executionState!.CancellationTokenSource != null && !executionState.CancellationTokenSource.Token.IsCancellationRequested)
        {
            executionState.CancellationTokenSource.Cancel();
        }

        _responseStore.UpdateTaskStatus(appChatId, appRunId, ResponseGenerationStatus.Cancelled, "Cancelled by user request");
        
        await Task.CompletedTask;
    }

    /// <summary>
    /// Converts application chat message to Microsoft.Extensions.AI ChatMessage format.
    /// </summary>
    /// <param name="message">Application-specific chat message.</param>
    /// <returns>ChatMessage object for the AI client.</returns>
    private static AIChatMessage ConvertToChatMessages(ChatMessageBase message)
    {
        var content = message switch
            {
                UserPrompt userPrompt => userPrompt.Content,
                BotResponse botResponse => botResponse.Content,
                _ => null
            };

            if (string.IsNullOrEmpty(content)){
                throw new InvalidOperationException("User's message content is empty.");
            }

            var role = message.Role switch
            {
                ChatMessageRole.User => ChatRole.User,
                ChatMessageRole.Bot => ChatRole.Assistant,
                ChatMessageRole.System => ChatRole.System,
                _ => ChatRole.User
            };

        return new AIChatMessage(role, content);
    }
} 

[stephentoub]

I just wanted to thank you all so much for this update!

You're welcome :)

I am curious to know if I am using, the OpenAIAssistantChatClient : IChatClient correctly?

It should be usable like any other IChatClient. The main thing to ensure, which it looks like you are, is that you're providing the thread ID as the ConversationId, as outlined in https://learn.microsoft.com/en-us/dotnet/ai/microsoft-extensions-ai#stateless-vs-stateful-clients.