GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

openai/openai-dotnet

Public

mirrored from https://github.com/openai/openai-dotnetAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
achandmsft-patch-2

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openai-dotnet/codegen/generator/src/Visitors

codegen/generator/src/Visitors/PaginationVisitor.cs

257lines · modecode

RawDownload
Latest commit unavailable.
1using System;
2using System.Collections.Generic;
3using System.Linq;
4using Microsoft.TypeSpec.Generator.ClientModel;
5using Microsoft.TypeSpec.Generator.Expressions;
6using Microsoft.TypeSpec.Generator.Primitives;
7using Microsoft.TypeSpec.Generator.Providers;
8using Microsoft.TypeSpec.Generator.Snippets;
9using Microsoft.TypeSpec.Generator.Statements;
10using static OpenAILibraryPlugin.Visitors.VisitorHelpers;
11
12namespace OpenAILibraryPlugin.Visitors;
13
14/// <summary>
15/// This visitor modifies GetRawPagesAsync methods to consider HasMore in addition to LastId when deciding whether to continue pagination.
16/// It also replaces specific parameters with an options type for pagination methods.
17/// </summary>
18public class PaginationVisitor : ScmLibraryVisitor
19{
20
21 private static readonly string[] _chatParamsToReplace = ["after", "before", "limit", "order", "model", "metadata"];
22 private static readonly Dictionary<string, string> _paramReplacementMap = new()
23 {
24 { "after", "AfterId" },
25 { "before", "LastId" },
26 { "limit", "PageSizeLimit" },
27 { "order", "Order" },
28 { "model", "Model" },
29 { "metadata", "Metadata" }
30 };
31 private static readonly Dictionary<string, (string ReturnType, string OptionsType, string[] ParamsToReplace)> _optionsReplacements = new()
32 {
33 {
34 "GetChatCompletions",
35 ("ChatCompletion", "ChatCompletionCollectionOptions", _chatParamsToReplace)
36 },
37 {
38 "GetChatCompletionsAsync",
39 ("ChatCompletion", "ChatCompletionCollectionOptions", _chatParamsToReplace)
40 }
41 };
42
43 protected override MethodProvider? VisitMethod(MethodProvider method)
44 {
45 // Try to handle pagination methods with options replacement
46 if (TryHandlePaginationMethodWithOptions(method))
47 {
48 return method;
49 }
50
51 // Try to handle GetRawPagesAsync methods for hasMore checks
52 if (TryHandleGetRawPagesAsyncMethod(method))
53 {
54 return method;
55 }
56
57 return base.VisitMethod(method);
58 }
59
60 /// <summary>
61 /// Handles pagination methods that need their parameters replaced with an options type.
62 /// </summary>
63 /// <param name="method">The method to potentially handle. Will be modified in place if handling is successful.</param>
64 /// <returns>True if the method was handled, false otherwise.</returns>
65 private bool TryHandlePaginationMethodWithOptions(MethodProvider method)
66 {
67 // Check if the method is one of the pagination methods we want to modify.
68 // If so, we will update its parameters to replace the specified parameters with the options type.
69 if (method.Signature.ReturnType is not null &&
70 method.Signature.ReturnType.Name.EndsWith("CollectionResult") &&
71 _optionsReplacements.TryGetValue(method.Signature.Name, out var options) &&
72 method.Signature.ReturnType.IsGenericType &&
73 method.Signature.ReturnType.Arguments.Count == 1 &&
74 method.Signature.ReturnType.Arguments[0].Name == options.ReturnType)
75 {
76 var optionsType = OpenAILibraryGenerator.Instance.OutputLibrary.TypeProviders.SingleOrDefault(t => t.Type.Name == options.OptionsType);
77 if (optionsType is not null)
78 {
79 // replace the method parameters with names in the _paramsToReplace array with the optionsType
80 var methodSignature = method.Signature;
81 var newParameters = methodSignature.Parameters.ToList();
82 int lastRemovedIndex = -1;
83 for (int i = 0; i < newParameters.Count; i++)
84 {
85 if (_chatParamsToReplace.Contains(newParameters[i].Name))
86 {
87 newParameters.RemoveAt(i);
88 lastRemovedIndex = i;
89 i--;
90 }
91 }
92 if (lastRemovedIndex >= 0)
93 {
94 newParameters.Insert(
95 lastRemovedIndex,
96 new ParameterProvider("options", $"The pagination options", optionsType.Type, defaultValue: Snippet.Default));
97
98 var newSignature = new MethodSignature(
99 methodSignature.Name,
100 methodSignature.Description,
101 methodSignature.Modifiers,
102 methodSignature.ReturnType,
103 methodSignature.ReturnDescription,
104 newParameters,
105 methodSignature.Attributes,
106 methodSignature.GenericArguments,
107 methodSignature.GenericParameterConstraints,
108 methodSignature.ExplicitInterface,
109 methodSignature.NonDocumentComment);
110
111 var optionsParam = newParameters[lastRemovedIndex];
112
113 // Update the method body statements to replace the old parameters with the new options parameter.
114 var statements = method.BodyStatements?.ToList() ?? new List<MethodBodyStatement>();
115 VisitExplodedMethodBodyStatements(statements!,
116 statement =>
117 {
118 // Check if the statement is a return statement
119 if (statement is ExpressionStatement exp && exp.Expression is KeywordExpression keyword && keyword.Keyword == "return")
120 {
121 // If it is, we will replace the parameters with the options parameter.
122 if (keyword.Expression is NewInstanceExpression newInstance &&
123 newInstance.Parameters.Count > 0)
124 {
125 // Create the new parameters with the options parameter.
126 var newParameters = new List<ValueExpression>();
127 foreach (var param in newInstance.Parameters)
128 {
129 if (param is VariableExpression varExpr && options.ParamsToReplace.Contains(varExpr.Declaration.RequestedName))
130 {
131 // Replace the parameter with the options parameter.
132 if (_paramReplacementMap.TryGetValue(varExpr.Declaration.RequestedName, out var replacement))
133 {
134 newParameters.Add(optionsParam.NullConditional().Property(replacement));
135 }
136 }
137 else if (param is InvokeMethodExpression invokeMethod && invokeMethod.MethodName == "ToString" &&
138 invokeMethod.InstanceReference is NullConditionalExpression nullConditional &&
139 nullConditional.Inner is VariableExpression varExpr2 &&
140 options.ParamsToReplace.Contains(varExpr2.Declaration.RequestedName))
141 {
142 // Replace the parameter with the options parameter.
143 if (_paramReplacementMap.TryGetValue(varExpr2.Declaration.RequestedName, out var replacement))
144 {
145 newParameters.Add(optionsParam.NullConditional().Property(replacement).NullConditional().Invoke("ToString", Array.Empty<ValueExpression>()));
146 }
147 }
148 else
149 {
150 // Keep the original parameter.
151 newParameters.Add(param);
152 }
153 }
154 // Create a new ExpressionStatement with the same children as the original, but with the new parameters.
155 return Snippet.Return(Snippet.New.Instance(newInstance.Type!, newParameters));
156 }
157 }
158 return statement;
159 });
160
161 method.Update(signature: newSignature, bodyStatements: statements);
162 return true;
163 }
164 }
165 }
166
167 return false;
168 }
169
170 /// <summary>
171 /// Handles GetRawPagesAsync methods to add hasMore == false checks for pagination.
172 /// </summary>
173 /// <param name="method">The method to potentially handle. Will be modified in place if handling is successful.</param>
174 /// <returns>True if the method was handled, false otherwise.</returns>
175 private bool TryHandleGetRawPagesAsyncMethod(MethodProvider method)
176 {
177 // If the method is GetRawPagesAsync and is internal, we will modify the body statements to add a check for hasMore == false.
178 // 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))
180 {
181 var statements = method.BodyStatements?.ToList() ?? new List<MethodBodyStatement>();
182 VisitExplodedMethodBodyStatements(
183 statements!,
184 statement =>
185 {
186 if (statement is IfStatement)
187 {
188 return GetUpdatedIfStatement(
189 statement,
190 expression =>
191 {
192 // Check if this is a binary expression with "==" operator
193 if (expression is ScopedApi scopedApi && scopedApi.Original is BinaryOperatorExpression binaryExpr && binaryExpr.Operator == "==")
194 {
195 // Check if left side is "nextToken" and right side is "null"
196 if (binaryExpr.Left is VariableExpression leftVar &&
197 leftVar.Declaration.RequestedName == "nextToken" &&
198 binaryExpr.Right is KeywordExpression rightKeyword &&
199 rightKeyword.Keyword == "null")
200 {
201 // Create "hasMore == null" condition
202 var hasMoreNullCheck = new BinaryOperatorExpression(
203 "==",
204 new MemberExpression(null, "hasMore"),
205 Snippet.False);
206
207 // Return "nextToken == null || hasMore == null"
208 return new BinaryOperatorExpression("||", binaryExpr, hasMoreNullCheck);
209 }
210 }
211 return expression;
212 },
213 "Plugin customization: add hasMore == false check to pagination condition");
214 }
215 else if (statement is WhileStatement whileStatement)
216 {
217 var statementList = whileStatement.Body
218 .SelectMany(bodyStatement => bodyStatement)
219 .ToList();
220
221 // Check for the assignment of nextToken and add hasMore assignment
222 for (int i = 0; i < statementList.Count; i++)
223 {
224 if (statementList[i] is ExpressionStatement expressionStatement &&
225 expressionStatement.Expression is AssignmentExpression assignmentExpression &&
226 assignmentExpression.Variable is VariableExpression variableExpression &&
227 variableExpression.Declaration.RequestedName == "nextToken" &&
228 assignmentExpression.Value is MemberExpression memberExpression &&
229 memberExpression.MemberName == "LastId")
230 {
231 // Create a new assignment for hasMore
232 var hasMoreAssignment = new AssignmentExpression(
233 new DeclarationExpression(typeof(bool), "hasMore"),
234 new MemberExpression(memberExpression.Inner, "HasMore"));
235
236 // Insert the new assignment before the existing one
237 statementList.Insert(i, hasMoreAssignment.Terminate());
238 statementList.Insert(i, new SingleLineCommentStatement("Plugin customization: add hasMore assignment"));
239 var updatedWhileStatement = new WhileStatement(whileStatement.Condition);
240 foreach (MethodBodyStatement bodyStatement in statementList)
241 {
242 updatedWhileStatement.Add(bodyStatement);
243 }
244 return updatedWhileStatement;
245 }
246 }
247 }
248 return statement;
249 });
250
251 method.Update(bodyStatements: statements);
252 return true;
253 }
254
255 return false;
256 }
257}