Merge pull request #1871 from solliancenet/cj-long-running-fixes

Improve status reporting for long-running completions

src/dotnet/Common/Models/Orchestration/LongRunningOperation.cs

@@ -51,6 +51,12 @@ public class LongRunningOperation
  [JsonPropertyName("ttl")]
  public int TTL { get; set; } = 604800;
 
+ /// <summary>
+ /// The number of tokens in the prompt.
+ /// </summary>
+ [JsonPropertyName("prompt_tokens")]
+ public int PromptTokens { get; set; } = 0;
+
  /// <summary>
  /// The result of the operation.
  /// </summary>

src/dotnet/Core/Interfaces/ICoreService.cs

@@ -2,7 +2,6 @@
 using FoundationaLLM.Common.Models.Conversation;
 using FoundationaLLM.Common.Models.Orchestration;
 using FoundationaLLM.Common.Models.Orchestration.Request;
-using FoundationaLLM.Common.Models.Orchestration.Response;
 using FoundationaLLM.Common.Models.ResourceProviders;
 using FoundationaLLM.Common.Models.ResourceProviders.Attachment;
 using FoundationaLLM.Common.Models.ResourceProviders.Configuration;
@@ -97,8 +96,8 @@ public interface ICoreService
  /// </summary>
  /// <param name="instanceId">The FoundationaLLM instance id.</param>
  /// <param name="operationId">The OperationId for which to retrieve the status.</param>
- /// <returns>Returns a <see cref="Message"/> object containing the agent's response message.</returns>
- Task<Message> GetCompletionOperationStatus(string instanceId, string operationId);
+ /// <returns>Returns a <see cref="LongRunningOperation"/> object containing the OperationId, Status, and result.</returns>
+ Task<LongRunningOperation> GetCompletionOperationStatus(string instanceId, string operationId);
 
  /// <summary>
  /// Uploads an attachment.

src/dotnet/Core/Services/CoreService.cs

@@ -7,6 +7,7 @@
 using FoundationaLLM.Common.Extensions;
 using FoundationaLLM.Common.Interfaces;
 using FoundationaLLM.Common.Models.Authentication;
+using FoundationaLLM.Common.Models.Azure.CosmosDB;
 using FoundationaLLM.Common.Models.Configuration.Branding;
 using FoundationaLLM.Common.Models.Conversation;
 using FoundationaLLM.Common.Models.Orchestration;
@@ -23,16 +24,16 @@
 using FoundationaLLM.Core.Interfaces;
 using FoundationaLLM.Core.Models.Configuration;
 using FoundationaLLM.Core.Utils;
+using Microsoft.AspNetCore.Http.HttpResults;
+using Microsoft.Azure.Cosmos.Serialization.HybridRow;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 using System.Text.Json;
 using System.Text.RegularExpressions;
-using FoundationaLLM.Common.Models.Azure.CosmosDB;
 using Conversation = FoundationaLLM.Common.Models.Conversation.Conversation;
 using LongRunningOperation = FoundationaLLM.Common.Models.Orchestration.LongRunningOperation;
 using Message = FoundationaLLM.Common.Models.Conversation.Message;
-using FoundationaLLM.Common.Models.Orchestration;
 
 namespace FoundationaLLM.Core.Services;
 
@@ -337,6 +338,38 @@ await _cosmosDBService.UpsertLongRunningOperationContextAsync(new LongRunningOpe
  // Start the completion operation.
  var result = await GetDownstreamAPIService(agentOption).StartCompletionOperation(instanceId, completionRequest);
 
+ if (result.Status == OperationStatus.Failed)
+ {
+ // In case the completion operation fails to start properly, we need to update the user and agent messages accordingly.
+
+ var patchOperations = new List<IPatchOperationItem>
+ {
+ new PatchOperationItem<Message>
+ {
+ ItemId = conversationItems.UserMessage.Id,
+ PropertyValues = new Dictionary<string, object?>
+ {
+ { "/status", OperationStatus.Failed },
+ { "/text", result.StatusMessage }
+ }
+ },
+ new PatchOperationItem<Message>
+ {
+ ItemId = conversationItems.AgentMessage.Id,
+ PropertyValues = new Dictionary<string, object?>
+ {
+ { "/status", OperationStatus.Failed },
+ { "/text", result.StatusMessage }
+ }
+ }
+ };
+
+ var patchedItems = await _cosmosDBService.PatchMultipleSessionsItemsInTransactionAsync(
+ completionRequest.SessionId!,
+ patchOperations
+ );
+ }
+
  return result;
  }
  catch (Exception ex)
@@ -356,7 +389,7 @@ await _cosmosDBService.UpsertLongRunningOperationContextAsync(new LongRunningOpe
  }
 
  /// <inheritdoc/>
- public async Task<Message> GetCompletionOperationStatus(string instanceId, string operationId)
+ public async Task<LongRunningOperation> GetCompletionOperationStatus(string instanceId, string operationId)
  {
  try
  {
@@ -368,20 +401,44 @@ public async Task<Message> GetCompletionOperationStatus(string instanceId, strin
  {
  // We've hit the hard stop time for the operation.
 
- await _cosmosDBService.PatchSessionsItemPropertiesAsync<Message>(
- operationContext.AgentMessageId,
- operationContext.SessionId,
- new Dictionary<string, object?>
+ var patchOperations = new List<IPatchOperationItem>
+ {
+ new PatchOperationItem<Message>
  {
- { "/status", OperationStatus.Failed },
- { "/text", "The completion operation has exceeded the maximum time allowed." }
- });
+ ItemId = operationContext.UserMessageId,
+ PropertyValues = new Dictionary<string, object?>
+ {
+ { "/status", OperationStatus.Failed },
+ { "/text", "The completion operation has exceeded the maximum time allowed." }
+ }
+ },
+ new PatchOperationItem<Message>
+ {
+ ItemId = operationContext.AgentMessageId,
+ PropertyValues = new Dictionary<string, object?>
+ {
+ { "/status", OperationStatus.Failed },
+ { "/text", "The completion operation has exceeded the maximum time allowed." }
+ }
+ }
+ };
 
- return new Message
+ var patchedItems = await _cosmosDBService.PatchMultipleSessionsItemsInTransactionAsync(
+ operationContext.SessionId,
+ patchOperations
+ );
+
+ return new LongRunningOperation
  {
  OperationId = operationId,
- Status = OperationStatus.Failed,
- Text = "The completion operation has exceeded the maximum time allowed."
+ StatusMessage = "The completion operation has exceeded the maximum time allowed.",
+ Result = new Message
+ {
+ OperationId = operationId,
+ Status = OperationStatus.Failed,
+ Text = "The completion operation has exceeded the maximum time allowed."
+ },
+ Status = OperationStatus.Failed
  };
  }
 
@@ -402,25 +459,39 @@ await _cosmosDBService.PatchSessionsItemPropertiesAsync<Message>(
  }
  }
 
- return agentMessage;
+ operationStatus.Result = agentMessage;
+ if (completionResponse != null)
+ {
+ operationStatus.PromptTokens = completionResponse.PromptTokens;
+ }
+
+ return operationStatus;
  }
 
- return new Message
+ operationStatus.Result = new Message
  {
  OperationId = operationId,
  Status = operationStatus.Status,
  Text = operationStatus.StatusMessage ?? "The completion operation is in progress."
  };
+
+ return operationStatus;
  }
  catch (Exception ex)
  {
  _logger.LogError(ex, "Error retrieving the status for the operation with id {OperationId}.",
  operationId);
- return new Message
+ return new LongRunningOperation
  {
  OperationId = operationId,
- Status = OperationStatus.Failed,
- Text = "Could not retrieve the status of the operation due to an internal error."
+ StatusMessage = "Could not retrieve the status of the operation due to an internal error.",
+ Result = new Message
+ {
+ OperationId = operationId,
+ Status = OperationStatus.Failed,
+ Text = "Could not retrieve the status of the operation due to an internal error."
+ },
+ Status = OperationStatus.Failed
  };
  }
  }
@@ -721,7 +792,11 @@ private async Task<Message> ProcessCompletionResponse(
  new PatchOperationItem<Message>
  {
  ItemId = operationContext.UserMessageId,
- PropertyValues = new Dictionary<string, object?> { { "/tokens", completionResponse.PromptTokens } }
+ PropertyValues = new Dictionary<string, object?>
+ {
+ { "/tokens", completionResponse.PromptTokens },
+ { "/status", operationStatus }
+ }
  },
  new PatchOperationItem<Message>
  {

src/dotnet/CoreAPI/Controllers/CompletionsController.cs

@@ -86,9 +86,9 @@ public async Task<ActionResult<LongRunningOperation>> StartCompletionOperation(s
  /// </summary>
  /// <param name="instanceId">The FoundationaLLM instance id.</param>
  /// <param name="operationId">The OperationId for which to retrieve the status.</param>
- /// <returns>Returns an <see cref="LongRunningOperation"/> object containing the OperationId and Status.</returns>
+ /// <returns>Returns a <see cref="LongRunningOperation"/> object containing the OperationId, Status, and result.</returns>
  [HttpGet("async-completions/{operationId}/status")]
- public async Task<Message> GetCompletionOperationStatus(string instanceId, string operationId) =>
+ public async Task<LongRunningOperation> GetCompletionOperationStatus(string instanceId, string operationId) =>
  await _coreService.GetCompletionOperationStatus(instanceId, operationId);
 
  /// <summary>

src/ui/UserPortal/components/ChatMessage.vue

@@ -313,6 +313,18 @@ export default {
 
  computed: {
  messageContent() {
+ if (this.message.status === 'Failed') {
+ const failedMessage = this.message.text ?? 'Failed to generate a response.';
+ return [
+ {
+ type: 'text',
+ content: failedMessage,
+ value: failedMessage,
+ origValue: failedMessage,
+ }
+ ];
+ }
+
  return this.message.content ?? [];
  },
 

src/ui/UserPortal/js/api.ts

@@ -1,6 +1,7 @@
 import type {
  Message,
  Session,
+ LongRunningOperation,
  UserProfile,
  CoreConfiguration,
  OneDriveWorkSchool,
@@ -148,8 +149,7 @@ export default {
  try {
  const response = await this.fetch(
  `/instances/${this.instanceId}/async-completions/${operationId}/status`,
- );
-
+ ) as LongRunningOperation;
  return response;
  } catch (error) {
  throw new Error(formatError(error));

src/ui/UserPortal/js/types/index.ts

@@ -81,6 +81,18 @@ export interface ChatSessionProperties {
  name: string;
 }
 
+export interface LongRunningOperation {
+ id?: string;
+ type: string;
+ operation_id?: string;
+ status: string;
+ status_message?: string;
+ last_updated?: Date;
+ ttl: number;
+ prompt_tokens: number;
+ result?: Message;
+}
+
 export interface CompletionPrompt {
  id: string;
  type: string;

src/ui/UserPortal/stores/appStore.ts

@@ -384,12 +384,18 @@ export const useAppStore = defineStore('app', {
 
  this.pollingInterval = setInterval(async () => {
  try {
- const updatedMessage = await api.checkProcessStatus(message.operation_id);
+ const statusResponse = await api.checkProcessStatus(message.operation_id);
+ const updatedMessage = statusResponse.result ?? {};
  this.currentMessages[this.currentMessages.length - 1] = {
  ...updatedMessage,
  renderId: this.currentMessages[this.currentMessages.length - 1].renderId,
  };
 
+ const userMessage = this.currentMessages[this.currentMessages.length - 2];
+ if (userMessage && statusResponse.prompt_tokens && userMessage.tokens !== statusResponse.prompt_tokens) {
+ userMessage.tokens = statusResponse.prompt_tokens;
+ }
+
  if (updatedMessage.status === 'Completed' || updatedMessage.status === 'Failed') {
  this.stopPolling();
  }