Skip to content

Commit f9d062b

Browse files
christotheschristothes
committed
Generate GetChatCompletionMessages with pagination
1 parent 27ac530 commit f9d062b

21 files changed

+1078
-46
lines changed

‎codegen/generator/src/Visitors/PaginationVisitor.cs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
namespace OpenAILibraryPlugin.Visitors;
1313

1414
/// <summary>
15-
/// This visitor modifies GetRawPagesAsync methods to consider HasMore in addition to LastId when deciding whether to continue pagination.
15+
/// This visitor modifies GetRawPagesAsync and GetRawPages methods to consider HasMore in addition to LastId when deciding whether to continue pagination.
1616
/// It also replaces specific parameters with an options type for pagination methods.
1717
/// </summary>
1818
public class PaginationVisitor : ScmLibraryVisitor
@@ -37,6 +37,14 @@ public class PaginationVisitor : ScmLibraryVisitor
3737
{
3838
"GetChatCompletionsAsync",
3939
("ChatCompletion", "ChatCompletionCollectionOptions", _chatParamsToReplace)
40+
},
41+
{
42+
"GetChatCompletionMessages",
43+
("ChatCompletionMessageListDatum", "ChatCompletionCollectionOptions", _chatParamsToReplace)
44+
},
45+
{
46+
"GetChatCompletionMessagesAsync",
47+
("ChatCompletionMessageListDatum", "ChatCompletionMessageCollectionOptions", _chatParamsToReplace)
4048
}
4149
};
4250

@@ -174,9 +182,9 @@ nullConditional.Inner is VariableExpression varExpr2 &&
174182
/// <returns>True if the method was handled, false otherwise.</returns>
175183
private bool TryHandleGetRawPagesAsyncMethod(MethodProvider method)
176184
{
177-
// If the method is GetRawPagesAsync and is internal, we will modify the body statements to add a check for hasMore == false.
185+
// If the method is GetRawPagesAsync or GetRawPages and is internal, we will modify the body statements to add a check for hasMore == false.
178186
// This is to ensure that pagination stops when hasMore is false, in addition to checking LastId.
179-
if (method.Signature.Name == "GetRawPagesAsync" && method.EnclosingType.DeclarationModifiers.HasFlag(TypeSignatureModifiers.Internal))
187+
if ((method.Signature.Name == "GetRawPagesAsync" || method.Signature.Name == "GetRawPages") && method.EnclosingType.DeclarationModifiers.HasFlag(TypeSignatureModifiers.Internal))
180188
{
181189
var statements = method.BodyStatements?.ToList() ?? new List<MethodBodyStatement>();
182190
VisitExplodedMethodBodyStatements(

‎specification/client/models/chat.models.tsp‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,26 @@ model ChatCompletionCollectionOptions {
3131
@query `model`?: string,
3232
}
3333

34+
alias ChatCompletionMessageCollectionOrderQueryParameter = {
35+
/**
36+
* Sort order by the `created_at` timestamp of the objects. `asc` for ascending order and`desc`
37+
* for descending order.
38+
*/
39+
@query order?: ChatCompletionCollectionOrder;
40+
};
41+
42+
union ChatCompletionMessageCollectionOrder {
43+
string,
44+
Ascending: "asc",
45+
Descending: "desc",
46+
}
47+
48+
@access(Access.public)
49+
@usage(Usage.input)
50+
model ChatCompletionMessageCollectionOptions {
51+
...CollectionAfterQueryParameter,
52+
...CollectionLimitQueryParameter,
53+
...ChatCompletionMessageCollectionOrderQueryParameter,
54+
}
55+
3456

‎src/Custom/Chat/ChatClient.Protocol.cs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@
77
namespace OpenAI.Chat;
88

99
/// <summary> The service client for the OpenAI Chat Completions endpoint. </summary>
10-
[CodeGenSuppress("GetChatCompletionMessagesAsync", typeof(string), typeof(string), typeof(int?), typeof(string), typeof(RequestOptions))]
11-
[CodeGenSuppress("GetChatCompletionMessages", typeof(string), typeof(string), typeof(int?), typeof(string), typeof(RequestOptions))]
1210
[CodeGenSuppress("UpdateChatCompletionAsync", typeof(string), typeof(BinaryContent), typeof(RequestOptions))]
1311
[CodeGenSuppress("UpdateChatCompletion", typeof(string), typeof(BinaryContent), typeof(RequestOptions))]
1412
public partial class ChatClient

‎src/Custom/Chat/ChatClient.cs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,6 @@ namespace OpenAI.Chat;
2020
[CodeGenSuppress("ChatClient", typeof(ClientPipeline), typeof(Uri))]
2121
[CodeGenSuppress("CompleteChat", typeof(ChatCompletionOptions), typeof(CancellationToken))]
2222
[CodeGenSuppress("CompleteChatAsync", typeof(ChatCompletionOptions), typeof(CancellationToken))]
23-
[CodeGenSuppress("GetChatCompletionMessages", typeof(string), typeof(string), typeof(int?), typeof(OpenAI.VectorStores.VectorStoreCollectionOrder?), typeof(CancellationToken))]
24-
[CodeGenSuppress("GetChatCompletionMessagesAsync", typeof(string), typeof(string), typeof(int?), typeof(OpenAI.VectorStores.VectorStoreCollectionOrder?), typeof(CancellationToken))]
2523
[CodeGenSuppress("UpdateChatCompletion", typeof(string), typeof(IDictionary<string, string>), typeof(CancellationToken))]
2624
[CodeGenSuppress("UpdateChatCompletionAsync", typeof(string), typeof(IDictionary<string, string>), typeof(CancellationToken))]
2725
public partial class ChatClient

‎src/Custom/Chat/ChatCompletionMessageCollectionOptions.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
2+
namespace OpenAI.Chat;
3+
4+
// CUSTOM: Use the correct namespace.
5+
[CodeGenType("ChatCompletionMessageCollectionOptions")]
6+
public partial class ChatCompletionMessageCollectionOptions {}

‎src/Custom/Chat/Internal/InternalChatCompletionMessageListDatum.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
namespace OpenAI.Chat;
22

33
[CodeGenType("ChatCompletionMessageListDatum")]
4-
internal partial class InternalChatCompletionMessageListDatum
4+
public partial class ChatCompletionMessageListDatum
55
{
66
// CUSTOM: Ensure enumerated value is used.
77
[CodeGenMember("Role")]
8-
internal ChatMessageRole Role { get; set; } = ChatMessageRole.Assistant;
8+
public ChatMessageRole Role { get; set; } = ChatMessageRole.Assistant;
99
}

‎src/Generated/ChatClient.cs‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,5 +90,61 @@ public virtual async Task<ClientResult> CompleteChatAsync(BinaryContent content,
9090
using PipelineMessage message = CreateCompleteChatRequest(content, options);
9191
return ClientResult.FromResponse(await Pipeline.ProcessMessageAsync(message, options).ConfigureAwait(false));
9292
}
93+
94+
[Experimental("OPENAI001")]
95+
public virtual CollectionResult GetChatCompletionMessages(string completionId, string after, int? limit, string order, RequestOptions options)
96+
{
97+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
98+
99+
return new ChatClientGetChatCompletionMessagesCollectionResult(
100+
this,
101+
completionId,
102+
after,
103+
limit,
104+
order,
105+
options);
106+
}
107+
108+
[Experimental("OPENAI001")]
109+
public virtual AsyncCollectionResult GetChatCompletionMessagesAsync(string completionId, string after, int? limit, string order, RequestOptions options)
110+
{
111+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
112+
113+
return new ChatClientGetChatCompletionMessagesAsyncCollectionResult(
114+
this,
115+
completionId,
116+
after,
117+
limit,
118+
order,
119+
options);
120+
}
121+
122+
[Experimental("OPENAI001")]
123+
public virtual CollectionResult<ChatCompletionMessageListDatum> GetChatCompletionMessages(string completionId, ChatCompletionCollectionOptions options = default, CancellationToken cancellationToken = default)
124+
{
125+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
126+
127+
return new ChatClientGetChatCompletionMessagesCollectionResultOfT(
128+
this,
129+
completionId,
130+
options?.AfterId,
131+
options?.PageSizeLimit,
132+
options?.Order?.ToString(),
133+
cancellationToken.CanBeCanceled ? new RequestOptions { CancellationToken = cancellationToken } : null);
134+
}
135+
136+
[Experimental("OPENAI001")]
137+
public virtual AsyncCollectionResult<ChatCompletionMessageListDatum> GetChatCompletionMessagesAsync(string completionId, ChatCompletionMessageCollectionOptions options = default, CancellationToken cancellationToken = default)
138+
{
139+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
140+
141+
return new ChatClientGetChatCompletionMessagesAsyncCollectionResultOfT(
142+
this,
143+
completionId,
144+
options?.AfterId,
145+
options?.PageSizeLimit,
146+
options?.Order?.ToString(),
147+
cancellationToken.CanBeCanceled ? new RequestOptions { CancellationToken = cancellationToken } : null);
148+
}
93149
}
94150
}

‎src/Generated/ChatClientGetChatCompletionMessagesAsyncCollectionResult.cs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
// <auto-generated/>
2+
3+
#nullable disable
4+
5+
using System;
6+
using System.ClientModel;
7+
using System.ClientModel.Primitives;
8+
using System.Collections.Generic;
9+
using OpenAI;
10+
11+
namespace OpenAI.Chat
12+
{
13+
internal partial class ChatClientGetChatCompletionMessagesAsyncCollectionResult : AsyncCollectionResult
14+
{
15+
private readonly ChatClient _client;
16+
private readonly string _completionId;
17+
private readonly string _after;
18+
private readonly int? _limit;
19+
private readonly string _order;
20+
private readonly RequestOptions _options;
21+
22+
public ChatClientGetChatCompletionMessagesAsyncCollectionResult(ChatClient client, string completionId, string after, int? limit, string order, RequestOptions options)
23+
{
24+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
25+
26+
_client = client;
27+
_completionId = completionId;
28+
_after = after;
29+
_limit = limit;
30+
_order = order;
31+
_options = options;
32+
}
33+
34+
public override async IAsyncEnumerable<ClientResult> GetRawPagesAsync()
35+
{
36+
PipelineMessage message = _client.CreateGetChatCompletionMessagesRequest(_completionId, _after, _limit, _order, _options);
37+
string nextToken = null;
38+
while (true)
39+
{
40+
ClientResult result = ClientResult.FromResponse(await _client.Pipeline.ProcessMessageAsync(message, _options).ConfigureAwait(false));
41+
yield return result;
42+
43+
// Plugin customization: add hasMore assignment
44+
bool hasMore = ((InternalChatCompletionMessageList)result).HasMore;
45+
nextToken = ((InternalChatCompletionMessageList)result).LastId;
46+
// Plugin customization: add hasMore == false check to pagination condition
47+
if (nextToken == null || hasMore == false)
48+
{
49+
yield break;
50+
}
51+
message = _client.CreateGetChatCompletionMessagesRequest(_completionId, nextToken, _limit, _order, _options);
52+
}
53+
}
54+
55+
public override ContinuationToken GetContinuationToken(ClientResult page)
56+
{
57+
string nextPage = ((InternalChatCompletionMessageList)page).LastId;
58+
if (nextPage != null)
59+
{
60+
return ContinuationToken.FromBytes(BinaryData.FromString(nextPage));
61+
}
62+
else
63+
{
64+
return null;
65+
}
66+
}
67+
}
68+
}

‎src/Generated/ChatClientGetChatCompletionMessagesAsyncCollectionResultOfT.cs‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
// <auto-generated/>
2+
3+
#nullable disable
4+
5+
using System;
6+
using System.ClientModel;
7+
using System.ClientModel.Primitives;
8+
using System.Collections.Generic;
9+
using System.Threading.Tasks;
10+
using OpenAI;
11+
12+
namespace OpenAI.Chat
13+
{
14+
internal partial class ChatClientGetChatCompletionMessagesAsyncCollectionResultOfT : AsyncCollectionResult<ChatCompletionMessageListDatum>
15+
{
16+
private readonly ChatClient _client;
17+
private readonly string _completionId;
18+
private readonly string _after;
19+
private readonly int? _limit;
20+
private readonly string _order;
21+
private readonly RequestOptions _options;
22+
23+
public ChatClientGetChatCompletionMessagesAsyncCollectionResultOfT(ChatClient client, string completionId, string after, int? limit, string order, RequestOptions options)
24+
{
25+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
26+
27+
_client = client;
28+
_completionId = completionId;
29+
_after = after;
30+
_limit = limit;
31+
_order = order;
32+
_options = options;
33+
}
34+
35+
public override async IAsyncEnumerable<ClientResult> GetRawPagesAsync()
36+
{
37+
PipelineMessage message = _client.CreateGetChatCompletionMessagesRequest(_completionId, _after, _limit, _order, _options);
38+
string nextToken = null;
39+
while (true)
40+
{
41+
ClientResult result = ClientResult.FromResponse(await _client.Pipeline.ProcessMessageAsync(message, _options).ConfigureAwait(false));
42+
yield return result;
43+
44+
// Plugin customization: add hasMore assignment
45+
bool hasMore = ((InternalChatCompletionMessageList)result).HasMore;
46+
nextToken = ((InternalChatCompletionMessageList)result).LastId;
47+
// Plugin customization: add hasMore == false check to pagination condition
48+
if (nextToken == null || hasMore == false)
49+
{
50+
yield break;
51+
}
52+
message = _client.CreateGetChatCompletionMessagesRequest(_completionId, nextToken, _limit, _order, _options);
53+
}
54+
}
55+
56+
public override ContinuationToken GetContinuationToken(ClientResult page)
57+
{
58+
string nextPage = ((InternalChatCompletionMessageList)page).LastId;
59+
if (nextPage != null)
60+
{
61+
return ContinuationToken.FromBytes(BinaryData.FromString(nextPage));
62+
}
63+
else
64+
{
65+
return null;
66+
}
67+
}
68+
69+
protected override async IAsyncEnumerable<ChatCompletionMessageListDatum> GetValuesFromPageAsync(ClientResult page)
70+
{
71+
foreach (ChatCompletionMessageListDatum item in ((InternalChatCompletionMessageList)page).Data)
72+
{
73+
yield return item;
74+
await Task.Yield();
75+
}
76+
}
77+
}
78+
}

‎src/Generated/ChatClientGetChatCompletionMessagesCollectionResult.cs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
// <auto-generated/>
2+
3+
#nullable disable
4+
5+
using System;
6+
using System.ClientModel;
7+
using System.ClientModel.Primitives;
8+
using System.Collections.Generic;
9+
using OpenAI;
10+
11+
namespace OpenAI.Chat
12+
{
13+
internal partial class ChatClientGetChatCompletionMessagesCollectionResult : CollectionResult
14+
{
15+
private readonly ChatClient _client;
16+
private readonly string _completionId;
17+
private readonly string _after;
18+
private readonly int? _limit;
19+
private readonly string _order;
20+
private readonly RequestOptions _options;
21+
22+
public ChatClientGetChatCompletionMessagesCollectionResult(ChatClient client, string completionId, string after, int? limit, string order, RequestOptions options)
23+
{
24+
Argument.AssertNotNullOrEmpty(completionId, nameof(completionId));
25+
26+
_client = client;
27+
_completionId = completionId;
28+
_after = after;
29+
_limit = limit;
30+
_order = order;
31+
_options = options;
32+
}
33+
34+
public override IEnumerable<ClientResult> GetRawPages()
35+
{
36+
PipelineMessage message = _client.CreateGetChatCompletionMessagesRequest(_completionId, _after, _limit, _order, _options);
37+
string nextToken = null;
38+
while (true)
39+
{
40+
ClientResult result = ClientResult.FromResponse(_client.Pipeline.ProcessMessage(message, _options));
41+
yield return result;
42+
43+
// Plugin customization: add hasMore assignment
44+
bool hasMore = ((InternalChatCompletionMessageList)result).HasMore;
45+
nextToken = ((InternalChatCompletionMessageList)result).LastId;
46+
// Plugin customization: add hasMore == false check to pagination condition
47+
if (nextToken == null || hasMore == false)
48+
{
49+
yield break;
50+
}
51+
message = _client.CreateGetChatCompletionMessagesRequest(_completionId, nextToken, _limit, _order, _options);
52+
}
53+
}
54+
55+
public override ContinuationToken GetContinuationToken(ClientResult page)
56+
{
57+
string nextPage = ((InternalChatCompletionMessageList)page).LastId;
58+
if (nextPage != null)
59+
{
60+
return ContinuationToken.FromBytes(BinaryData.FromString(nextPage));
61+
}
62+
else
63+
{
64+
return null;
65+
}
66+
}
67+
}
68+
}

0 commit comments

Comments
 (0)