Add token stats and raw completion info; prevent jailbreaking

src/typechat.meai/ChatLanguageModel.cs

@@ -37,13 +37,52 @@ public ChatLanguageModel(IChatClient chatClient, ModelInfo model)
     /// <param name="settings"></param>
     /// <param name="cancelToken"></param>
     /// <returns></returns>
-    public async Task<string> CompleteAsync(Prompt prompt, TranslationSettings? settings = null, CancellationToken cancelToken = default)
+    public async Task<LanguageModelResponse> CompleteAsync(Prompt prompt, TranslationSettings? settings = null, CancellationToken cancelToken = default)
     {
         List<ChatMessage> history = ToHistory(prompt);
         ChatOptions? options = ToRequestSettings(settings);
         var response = await _chatClient.GetResponseAsync(history, options: options, cancellationToken: cancelToken).ConfigureAwait(false);
 
-        return response.Text;
+        return new LanguageModelResponse(response.Text, BuildInfo(response));
+    }
+
+    private static CompletionInfo BuildInfo(ChatResponse response)
+    {
+        var info = new CompletionInfo
+        {
+            Model = response.ModelId,
+            FinishReason = MapFinishReason(response.FinishReason)
+        };
+        UsageDetails? usage = response.Usage;
+        if (usage is not null)
+        {
+            info.Usage = new TokenUsage(
+                (int)(usage.InputTokenCount ?? 0),
+                (int)(usage.OutputTokenCount ?? 0),
+                (int)(usage.TotalTokenCount ?? 0));
+        }
+        return info;
+    }
+
+    private static CompletionFinishReason? MapFinishReason(ChatFinishReason? reason)
+    {
+        if (reason is null)
+        {
+            return null;
+        }
+        switch (reason.Value.Value?.ToLowerInvariant())
+        {
+            case "stop":
+                return CompletionFinishReason.Stop;
+            case "length":
+                return CompletionFinishReason.Length;
+            case "content_filter":
+                return CompletionFinishReason.ContentFilter;
+            case "tool_calls":
+                return CompletionFinishReason.ToolCalls;
+            default:
+                return CompletionFinishReason.Other;
+        }
     }
 
     private List<ChatMessage> ToHistory(Prompt prompt)

src/typechat.sk/ChatLanguageModel.cs

@@ -51,7 +51,7 @@ public ChatLanguageModel(IChatCompletionService service, ModelInfo model)
     /// <param name="settings"></param>
     /// <param name="cancelToken"></param>
     /// <returns></returns>
-    public async Task<string> CompleteAsync(Prompt prompt, TranslationSettings? settings = null, CancellationToken cancelToken = default)
+    public async Task<LanguageModelResponse> CompleteAsync(Prompt prompt, TranslationSettings? settings = null, CancellationToken cancelToken = default)
     {
         ChatHistory history = ToHistory(prompt);
         OpenAIPromptExecutionSettings? requestSettings = ToRequestSettings(settings);

src/typechat.sk/TextCompletionModel.cs

@@ -40,7 +40,7 @@ public TextCompletionModel(ITextGenerationService service, ModelInfo model)
     /// </summary>
     public bool IncludeSectionSource { get; set; } = true;
 
-    public async Task<string> CompleteAsync(Prompt prompt, TranslationSettings? settings = null, CancellationToken cancelToken = default)
+    public async Task<LanguageModelResponse> CompleteAsync(Prompt prompt, TranslationSettings? settings = null, CancellationToken cancelToken = default)
     {
         OpenAIPromptExecutionSettings? requestSettings = ToRequestSettings(settings);
         string request = prompt.ToString(IncludeSectionSource);

src/typechat/CompletionFinishReason.cs

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+namespace Microsoft.TypeChat;
+
+/// <summary>
+/// Normalized reason a completion stopped generating, mapped from the underlying API so callers can
+/// detect conditions such as truncation (<see cref="Length"/>) or filtered output
+/// (<see cref="ContentFilter"/>) without special-casing each API variant.
+/// </summary>
+public enum CompletionFinishReason
+{
+    /// <summary>The model stopped naturally (Chat Completions "stop" / Responses "completed").</summary>
+    Stop,
+
+    /// <summary>Output was truncated by a token limit (Chat "length" / Responses "max_output_tokens").</summary>
+    Length,
+
+    /// <summary>Output was withheld or truncated by a content filter.</summary>
+    ContentFilter,
+
+    /// <summary>The model emitted tool/function calls.</summary>
+    ToolCalls,
+
+    /// <summary>A provider-specific reason that doesn't map to one of the common values.</summary>
+    Other,
+}

src/typechat/CompletionInfo.cs

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+
+namespace Microsoft.TypeChat;
+
+/// <summary>
+/// Metadata about a completion/translation. It is attached to <see cref="Result{T}.Info"/> on a
+/// successful result so callers can inspect details such as token usage (for example, to track cost
+/// or telemetry) without changing how existing results are consumed.
+/// </summary>
+public class CompletionInfo
+{
+    /// <summary>
+    /// The model that actually produced the completion, as reported by the API. This can differ from
+    /// the requested model (for example, when an alias resolves to a dated model version).
+    /// </summary>
+    public string? Model { get; set; }
+
+    /// <summary>
+    /// Normalized token usage for the request, when the API reports it.
+    /// </summary>
+    public TokenUsage? Usage { get; set; }
+
+    /// <summary>
+    /// Normalized reason the completion stopped, when the API reports it.
+    /// </summary>
+    public CompletionFinishReason? FinishReason { get; set; }
+
+    /// <summary>
+    /// Number of repair attempts TypeChat made before producing this result (0 = succeeded on the
+    /// first attempt). Populated by <see cref="JsonTranslator{T}"/>, not by the language model.
+    /// </summary>
+    public int? RepairAttempts { get; set; }
+
+    /// <summary>
+    /// The raw, unmodified response body returned by the API, when available. Use this to access
+    /// provider- or API-specific fields that are not normalized onto the other properties.
+    /// </summary>
+    public JsonElement? Raw { get; set; }
+}

src/typechat/HttpEx.cs

@@ -50,6 +50,51 @@ internal static async Task<Response> GetJsonResponseAsync<Request, Response>(thi
         return default;
     }
 
+    /// <summary>
+    /// Like <see cref="GetJsonResponseAsync"/>, but also returns the raw (unparsed) JSON response body
+    /// so callers can surface it (for example, on <see cref="CompletionInfo.Raw"/>).
+    /// </summary>
+    internal static async Task<(TResponse Response, string Raw)> GetJsonResponseWithRawAsync<TRequest, TResponse>(this HttpClient client, string endpoint, TRequest request, int maxRetries, int retryPauseMs, string? apiToken = null)
+    {
+        int retryCount = 0;
+        while (true)
+        {
+            HttpRequestMessage httpRequest = new HttpRequestMessage(HttpMethod.Post, endpoint)
+            {
+                Content = Json.ToJsonMessage(request)
+            };
+            try
+            {
+                if (!string.IsNullOrEmpty(apiToken))
+                {
+                    httpRequest.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiToken);
+                }
+                HttpResponseMessage response = await client.SendAsync(httpRequest).ConfigureAwait(false);
+                if (response.StatusCode == HttpStatusCode.OK)
+                {
+                    string raw = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
+                    return (Json.Parse<TResponse>(raw), raw);
+                }
+                if (!response.StatusCode.IsTransientError() || retryCount >= maxRetries)
+                {
+                    // Let HttpClient throw an exception
+                    response.EnsureSuccessStatusCode();
+                    break;
+                }
+                if (retryPauseMs > 0)
+                {
+                    await Task.Delay(retryPauseMs).ConfigureAwait(false);
+                }
+                retryCount++;
+            }
+            finally
+            {
+                httpRequest.Dispose();
+            }
+        }
+        return default;
+    }
+
     internal static bool IsTransientError(this HttpStatusCode status)
     {
         switch (status)

src/typechat/ILanguageModel.cs

@@ -19,6 +19,6 @@ public interface ILanguageModel
     /// <param name="prompt">prompt</param>
     /// <param name="settings">translation settings such as temperature</param>
     /// <param name="cancelToken">cancellation token</param>
-    /// <returns></returns>
-    Task<string> CompleteAsync(Prompt prompt, TranslationSettings? settings, CancellationToken cancelToken);
+    /// <returns>The completion text plus optional <see cref="CompletionInfo"/> metadata.</returns>
+    Task<LanguageModelResponse> CompleteAsync(Prompt prompt, TranslationSettings? settings, CancellationToken cancelToken);
 }

src/typechat/JsonTranslator.cs

@@ -187,6 +187,19 @@ public Task<T> TranslateAsync(string request, CancellationToken cancelToken = de
         return TranslateAsync(request, null, null, cancelToken);
     }
 
+    /// <summary>
+    /// Translate a natural language request into a <see cref="Result{T}"/> of type 'T' without
+    /// throwing on failure. On success, the result's <see cref="Result{T}.Info"/> carries completion
+    /// stats (token usage, finish reason, repair count, and the raw response when available).
+    /// </summary>
+    /// <param name="request">text request</param>
+    /// <param name="cancelToken">optional cancel token</param>
+    /// <returns>A <see cref="Result{T}"/> containing the object of type T, or a failure.</returns>
+    public Task<Result<T>> TranslateToResultAsync(string request, CancellationToken cancelToken = default)
+    {
+        return TranslateToResultAsync(request, null, null, cancelToken);
+    }
+
     /// <summary>
     /// Translate a natural language request into an object of type 'T'
     /// </summary>
@@ -196,12 +209,48 @@ public Task<T> TranslateAsync(string request, CancellationToken cancelToken = de
     /// <param name="cancelToken"></param>
     /// <returns>Result containing object of type T</returns>
     /// <exception cref="TypeChatException"></exception>
+    // TODO: This throwing overload will be deprecated in the future in favor of TranslateToResultAsync,
+    // which returns a Result<T> (carrying CompletionInfo stats) instead of throwing on failure.
     public async Task<T> TranslateAsync(
         Prompt request,
         IList<IPromptSection>? preamble,
         TranslationSettings? requestSettings = null,
         CancellationToken cancelToken = default
     )
+    {
+        Result<T> result = await TranslateWorkerAsync(request, preamble, requestSettings, throwOnFailure: true, cancelToken).ConfigureAwait(false);
+        return result.Value;
+    }
+
+    /// <summary>
+    /// Translate a natural language request into a <see cref="Result{T}"/> of type 'T'. Unlike
+    /// <see cref="TranslateAsync(Prompt, IList{IPromptSection}, TranslationSettings, CancellationToken)"/>,
+    /// this method does not throw on failure: it returns a failed <see cref="Result{T}"/> instead. On
+    /// success, the result's <see cref="Result{T}.Info"/> carries completion stats (token usage,
+    /// finish reason, repair count, and the raw response when available).
+    /// </summary>
+    /// <param name="request">natural language request</param>
+    /// <param name="preamble">optional preamble prepended to the prompt</param>
+    /// <param name="requestSettings">optional translation settings</param>
+    /// <param name="cancelToken">optional cancel token</param>
+    /// <returns>A <see cref="Result{T}"/> containing the object of type T, or a failure.</returns>
+    public Task<Result<T>> TranslateToResultAsync(
+        Prompt request,
+        IList<IPromptSection>? preamble = null,
+        TranslationSettings? requestSettings = null,
+        CancellationToken cancelToken = default
+    )
+    {
+        return TranslateWorkerAsync(request, preamble, requestSettings, throwOnFailure: false, cancelToken);
+    }
+
+    private async Task<Result<T>> TranslateWorkerAsync(
+        Prompt request,
+        IList<IPromptSection>? preamble,
+        TranslationSettings? requestSettings,
+        bool throwOnFailure,
+        CancellationToken cancelToken
+    )
     {
         ArgumentVerify.ThrowIfNull(request, nameof(request));
 
@@ -210,7 +259,8 @@ public async Task<T> TranslateAsync(
         int repairAttempts = 0;
         while (true)
         {
-            string responseText = await GetResponseAsync(prompt, requestSettings, cancelToken).ConfigureAwait(false);
+            LanguageModelResponse modelResponse = await GetResponseAsync(prompt, requestSettings, cancelToken).ConfigureAwait(false);
+            string responseText = modelResponse.Text;
 
             JsonResponse jsonResponse = JsonResponse.Parse(responseText);
             Result<T> validationResult;
@@ -219,6 +269,7 @@ public async Task<T> TranslateAsync(
                 validationResult = ValidateJson(jsonResponse.Json);
                 if (validationResult.Success)
                 {
+                    validationResult.Info = MergeInfo(modelResponse.Info, repairAttempts);
                     return validationResult;
                 }
             }
@@ -236,7 +287,11 @@ public async Task<T> TranslateAsync(
             ++repairAttempts;
             if (repairAttempts > _maxRepairAttempts)
             {
-                TypeChatException.ThrowJsonValidation(request, jsonResponse, validationResult.Message);
+                if (throwOnFailure)
+                {
+                    TypeChatException.ThrowJsonValidation(request, jsonResponse, validationResult.Message);
+                }
+                return validationResult;
             }
             NotifyEvent(AttemptingRepair, validationResult.Message);
 
@@ -251,17 +306,24 @@ public async Task<T> TranslateAsync(
         }
     }
 
+    private static CompletionInfo MergeInfo(CompletionInfo? info, int repairAttempts)
+    {
+        info ??= new CompletionInfo();
+        info.RepairAttempts = repairAttempts;
+        return info;
+    }
+
     protected virtual Prompt CreateRequestPrompt(Prompt request, IList<IPromptSection> preamble)
     {
         return _prompts.CreateRequestPrompt(_validator.Schema, request, preamble);
     }
 
-    protected virtual async Task<string> GetResponseAsync(Prompt prompt, TranslationSettings requestSettings, CancellationToken cancelToken)
+    protected virtual async Task<LanguageModelResponse> GetResponseAsync(Prompt prompt, TranslationSettings requestSettings, CancellationToken cancelToken)
     {
         NotifyEvent(SendingPrompt, prompt);
-        string responseText = await _model.CompleteAsync(prompt, requestSettings, cancelToken).ConfigureAwait(false);
-        NotifyEvent(CompletionReceived, responseText);
-        return responseText;
+        LanguageModelResponse response = await _model.CompleteAsync(prompt, requestSettings, cancelToken).ConfigureAwait(false);
+        NotifyEvent(CompletionReceived, response.Text);
+        return response;
     }
 
     protected virtual PromptSection CreateRepairPrompt(string responseText, Result<T> validationResult)

src/typechat/JsonTranslatorPrompts.cs

@@ -71,17 +71,17 @@ public static PromptSection IntroSection(string typeName, string schema)
     public static PromptSection RequestSection(string request)
     {
         PromptSection requestSection = new PromptSection();
-        requestSection += "The following is a user request:\n";
-        requestSection += $"\"\"\"\n{request}\n\"\"\"\n";
+        requestSection += "The following is a user request encoded as a JSON string:\n";
+        requestSection += $"{Json.Stringify(request, false)}\n";
         requestSection += "The following is the user request translated into a JSON object with 2 spaces of indentation and no properties with the value undefined:\n";
         return requestSection;
     }
 
     public static string RepairPrompt(string validationError)
     {
         validationError ??= string.Empty;
-        return "The JSON object is invalid for the following reason:\n" +
-               $"{validationError}\n" +
+        return "The JSON object is invalid. The following is the validation error encoded as a JSON string:\n" +
+               $"{Json.Stringify(validationError, false)}\n" +
                "The following is a revised JSON object. Do not include explanations.\n";
     }
 }