Skip to content

Commit 97df550

Browse files
mythzmythz
committed
Add support OnChatCompletionSuccess/FailedAsync callbacks + ChatCompletionLog data model
1 parent 37509ad commit 97df550

File tree

4 files changed

+168
-7
lines changed

4 files changed

+168
-7
lines changed

‎ServiceStack/src/ServiceStack.AI.Chat/ChatCompletionLog.cs‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
using ServiceStack.DataAnnotations;
2+
using ServiceStack.Web;
3+
4+
namespace ServiceStack.AI;
5+
6+
public class ChatCompletionLog : IMeta
7+
{
8+
[AutoIncrement]
9+
public long Id { get; set; }
10+
11+
/// <summary>
12+
/// Unique user-specified or system generated GUID for Job
13+
/// </summary>
14+
[Index(Unique = true)] public virtual string? RefId { get; set; }
15+
16+
/// <summary>
17+
/// Associate Job with a tag group
18+
/// </summary>
19+
public virtual string? Tag { get; set; }
20+
21+
/// <summary>
22+
/// The ASP .NET Identity Auth User Id to populate the IRequest Context ClaimsPrincipal and User Session
23+
/// </summary>
24+
public virtual string? UserId { get; set; }
25+
26+
/// <summary>
27+
/// The API Key, if one was used to access the Chat Service
28+
/// </summary>
29+
public virtual string? ApiKey { get; set; }
30+
31+
public string Model { get; set; }
32+
33+
public string? UserPrompt { get; set; }
34+
35+
public string? Answer { get; set; }
36+
37+
/// <summary>
38+
/// JSON Body of Request
39+
/// </summary>
40+
[StringLength(StringLengthAttribute.MaxText)]
41+
public virtual string RequestBody { get; set; }
42+
43+
/// <summary>
44+
/// The Response DTO JSON Body
45+
/// </summary>
46+
[StringLength(StringLengthAttribute.MaxText)]
47+
public virtual string? ResponseBody { get; set; }
48+
49+
public virtual string? ErrorCode { get; set; }
50+
51+
public virtual ResponseStatus? Error { get; set; }
52+
53+
[Index] public virtual DateTime CreatedDate { get; set; }
54+
55+
public virtual int? DurationMs { get; set; }
56+
57+
public int? PromptTokens { get; set; }
58+
59+
public int? CompletionTokens { get; set; }
60+
61+
public virtual Dictionary<string, string>? Meta { get; set; }
62+
}
63+
64+
public static class ChatCompletionLogUtils
65+
{
66+
public static ChatCompletionLog ToChatCompletionLog(this IRequest req, ChatCompletion request, ChatResponse response, string? refId = null)
67+
{
68+
ArgumentNullException.ThrowIfNull(request);
69+
ArgumentNullException.ThrowIfNull(response);
70+
71+
var userId = req.GetClaimsPrincipal()?.GetUserId();
72+
var apiKey = req.GetApiKey();
73+
userId ??= apiKey?.UserAuthId;
74+
userId ??= req.GetSession()?.UserAuthId;
75+
var duration = req.GetElapsed();
76+
77+
return new ChatCompletionLog
78+
{
79+
RefId = refId ?? req.GetTraceId() ?? Guid.NewGuid().ToString("N"),
80+
UserId = userId,
81+
ApiKey = apiKey?.Key,
82+
Model = request.Model,
83+
UserPrompt = request.GetUserPrompt(),
84+
Answer = response.GetAnswer(),
85+
RequestBody = request.ToJson(),
86+
ResponseBody = response.ToJson(),
87+
CreatedDate = DateTime.UtcNow,
88+
DurationMs = duration != TimeSpan.Zero ? (int)duration.TotalMilliseconds : null,
89+
PromptTokens = response.Usage?.PromptTokens,
90+
CompletionTokens = response.Usage?.CompletionTokens,
91+
};
92+
}
93+
94+
public static ChatCompletionLog ToChatCompletionLog(this IRequest req, ChatCompletion request, Exception ex, string? refId = null)
95+
{
96+
ArgumentNullException.ThrowIfNull(request);
97+
ArgumentNullException.ThrowIfNull(ex);
98+
99+
var userId = req.GetClaimsPrincipal()?.GetUserId();
100+
var apiKey = req.GetApiKey();
101+
userId ??= apiKey?.UserAuthId;
102+
userId ??= req.GetSession()?.UserAuthId;
103+
var duration = req.GetElapsed();
104+
105+
var status = ex.ToResponseStatus();
106+
107+
return new ChatCompletionLog
108+
{
109+
RefId = refId ?? req.GetTraceId() ?? Guid.NewGuid().ToString("N"),
110+
UserId = userId,
111+
ApiKey = apiKey?.Key,
112+
Model = request.Model,
113+
UserPrompt = request.GetUserPrompt(),
114+
RequestBody = request.ToJson(),
115+
ErrorCode = status.ErrorCode,
116+
Error = status,
117+
CreatedDate = DateTime.UtcNow,
118+
DurationMs = duration != TimeSpan.Zero ? (int)duration.TotalMilliseconds : null,
119+
};
120+
}
121+
}

‎ServiceStack/src/ServiceStack.AI.Chat/ChatFeature.cs‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -275,14 +275,22 @@ public async Task DisableProviderAsync(string requestId)
275275
}
276276

277277
public Func<ChatCompletion, IRequest, Task<ChatResponse>> ChatCompletionAsync { get; set; }
278-
278+
public Func<ChatCompletion, ChatResponse, IRequest, Task>? OnChatCompletionSuccessAsync { get; set; }
279+
public Func<ChatCompletion, Exception, IRequest, Task>? OnChatCompletionFailedAsync { get; set; }
280+
279281
public async Task<ChatResponse> DefaultChatCompletionAsync(ChatCompletion request, IRequest req)
280282
{
281283
var candidateProviders = Providers
282284
.Where(x => x.Value.Models.ContainsKey(request.Model))
283285
.ToList();
284286
if (candidateProviders.Count == 0)
285287
throw HttpError.NotFound($"Model {request.Model} not found");
288+
289+
var oLong = req.GetItem(Keywords.RequestDuration);
290+
if (oLong == null)
291+
{
292+
req.SetItem(Keywords.RequestDuration, System.Diagnostics.Stopwatch.GetTimestamp());
293+
}
286294

287295
Exception? firstEx = null;
288296
var i = 0;
@@ -295,6 +303,9 @@ public async Task<ChatResponse> DefaultChatCompletionAsync(ChatCompletion reques
295303
var provider = entry.Value;
296304
chatRequest.Model = request.Model;
297305
var ret = await provider.ChatAsync(chatRequest).ConfigAwait();
306+
var onSuccess = OnChatCompletionSuccessAsync;
307+
if (onSuccess != null)
308+
await onSuccess(chatRequest, ret, req).ConfigAwait();
298309
return ret;
299310
}
300311
catch (Exception ex)
@@ -304,10 +315,12 @@ public async Task<ChatResponse> DefaultChatCompletionAsync(ChatCompletion reques
304315
firstEx ??= ex;
305316
}
306317
}
307-
if (firstEx != null)
308-
throw firstEx;
309-
310-
throw HttpError.NotFound($"Model {request.Model} not found");
318+
319+
firstEx ??= HttpError.NotFound($"Model {request.Model} not found");
320+
var onFailed = OnChatCompletionFailedAsync;
321+
if (onFailed != null)
322+
await onFailed(request, firstEx, req).ConfigAwait();
323+
throw firstEx;
311324
}
312325
}
313326

‎ServiceStack/src/ServiceStack.AI.Chat/Message.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,18 @@ public static AiMessage File(string fileData, string? filename=null, string? tex
9292

9393
public static class MessageUtils
9494
{
95+
public static string? GetUserPrompt(this ChatCompletion request)
96+
{
97+
var textContents = request.Messages
98+
.Where(x => x.Role is "user" or null)
99+
.SelectMany(x => x.Content ?? [])
100+
.Where(x => x is AiTextContent)
101+
.Cast<AiTextContent>()
102+
.ToList();
103+
104+
return textContents.LastOrDefault()?.Text;
105+
}
106+
95107
public static string? GetAnswer(this ChatResponse? response)
96108
{
97109
var sb = StringBuilderCache.Allocate();

‎ServiceStack/tests/AdhocNew/Configure.AI.Chat.cs‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
using ServiceStack.AI;
2+
using ServiceStack.Data;
3+
using ServiceStack.OrmLite;
24

35
[assembly: HostingStartup(typeof(MyApp.ConfigureAiChat))]
46

@@ -8,10 +10,23 @@ public class ConfigureAiChat : IHostingStartup
810
{
911
public void Configure(IWebHostBuilder builder) => builder
1012
.ConfigureServices(services => {
11-
services.AddPlugin(new ChatFeature());
13+
services.AddPlugin(new ChatFeature
14+
{
15+
OnChatCompletionSuccessAsync = async (request, response, req) => {
16+
using var db = await req.Resolve<IDbConnectionFactory>().OpenAsync();
17+
await db.InsertAsync(req.ToChatCompletionLog(request, response));
18+
},
19+
OnChatCompletionFailedAsync = async (request, exception, req) => {
20+
using var db = await req.Resolve<IDbConnectionFactory>().OpenAsync();
21+
await db.InsertAsync(req.ToChatCompletionLog(request, exception));
22+
},
23+
});
1224

1325
services.ConfigurePlugin<MetadataFeature>(feature => {
1426
feature.AddPluginLink("/chat", "AI Chat");
1527
});
28+
}).ConfigureAppHost(appHost => {
29+
using var db = appHost.Resolve<IDbConnectionFactory>().Open();
30+
db.CreateTableIfNotExists<ChatCompletionLog>();
1631
});
17-
}
32+
}

0 commit comments

Comments
 (0)