1+ using System ;
2+ using System . Collections . Generic ;
3+ using System . Linq ;
4+ using Microsoft . TypeSpec . Generator . ClientModel ;
5+ using Microsoft . TypeSpec . Generator . Expressions ;
6+ using Microsoft . TypeSpec . Generator . Primitives ;
7+ using Microsoft . TypeSpec . Generator . Providers ;
8+ using Microsoft . TypeSpec . Generator . Snippets ;
9+ using Microsoft . TypeSpec . Generator . Statements ;
10+ using static OpenAILibraryPlugin . Visitors . VisitorHelpers ;
11+
12+ namespace 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>
18+ public class MetadataQueryParamVisitor : 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+ /// <summary>
44+ /// Visits Create*Request methods to modify how metadata query parameters are handled.
45+ /// It replaces the following statements:
46+ /// <code>
47+ /// List<object> list = new List<object>();
48+ /// foreach (var @param in metadata)
49+ /// {
50+ /// uri.AppendQuery($"metadata[{@param.Key}]", @param.Value, true);
51+ /// list.Add(@param.Key);
52+ /// list.Add(@param.Value);
53+ /// }
54+ /// uri.AppendQueryDelimited("metadata", list, ",", null, true);
55+ /// </code>
56+ /// with:
57+ /// <code>
58+ /// foreach (var @param in metadata)
59+ /// {
60+ /// uri.AppendQuery($"metadata[{@param.Key}]", @param.Value, true);
61+ /// }
62+ /// </summary>
63+ /// <param name="method"></param>
64+ /// <returns></returns>
65+ protected override MethodProvider ? VisitMethod ( MethodProvider method )
66+ {
67+ // Check if the method is one of the Create*Request methods and has a signature that takes a metadata parameter like IDictionary<string, string> metadata
68+ if ( method . Signature . Name . StartsWith ( "Create" ) && method . Signature . Name . EndsWith ( "Request" ) &&
69+ method . Signature . Parameters . Any ( p => p . Type . IsDictionary && p . Name == "metadata" ) )
70+ {
71+ ValueExpression ? uri = null ;
72+ var statements = method . BodyStatements ? . ToList ( ) ?? new List < MethodBodyStatement > ( ) ;
73+ VisitExplodedMethodBodyStatements (
74+ statements ! ,
75+ statement =>
76+ {
77+ // Check if the statement is an assignment to a variable named "uri"
78+ // Capture it if so
79+ if ( statement is ExpressionStatement expressionStatement &&
80+ expressionStatement . Expression is AssignmentExpression assignmentExpression &&
81+ assignmentExpression . Variable is DeclarationExpression declarationExpression &&
82+ declarationExpression . Variable is VariableExpression variableExpression &&
83+ variableExpression . Declaration . RequestedName == "uri" )
84+ {
85+ uri = variableExpression ;
86+ }
87+ // Try to remove the unnecessary list declaration
88+ if ( statement is ExpressionStatement expressionStatement2 &&
89+ expressionStatement2 . Expression is AssignmentExpression assignmentExpression2 &&
90+ assignmentExpression2 . Variable is DeclarationExpression declarationExpression2 &&
91+ declarationExpression2 . Variable is VariableExpression variableExpression2 &&
92+ variableExpression2 . Declaration . RequestedName == "list" &&
93+ variableExpression2 . Type . IsCollection && variableExpression2 . Type . IsGenericType )
94+ {
95+ // Remove the list declaration
96+ return new SingleLineCommentStatement ( "Plugin customization: remove unnecessary list declaration" ) ;
97+ }
98+
99+ if ( uri is not null &&
100+ statement is ForEachStatement foreachStatement &&
101+ foreachStatement . Enumerable is DictionaryExpression dictionaryExpression &&
102+ dictionaryExpression . Original is VariableExpression variable &&
103+ variable . Declaration . RequestedName == "metadata" )
104+ {
105+ var formatString = new FormattableStringExpression ( "metadata[{0}]" , [ foreachStatement . ItemVariable . Property ( "Key" ) ] ) ;
106+ var appendQueryStatement = uri . Invoke ( "AppendQuery" , [ formatString , foreachStatement . ItemVariable . Property ( "Value" ) , new KeywordExpression ( "true" , null ) ] ) ;
107+ foreachStatement . Body . Clear ( ) ;
108+ foreachStatement . Body . Add ( new SingleLineCommentStatement ( "Plugin customization: Properly handle metadata query parameters" ) ) ;
109+ foreachStatement . Body . Add ( new ExpressionStatement ( appendQueryStatement ) ) ;
110+ }
111+
112+ // Remove the call to AppendQueryDelimited for metadata
113+ if ( statement is ExpressionStatement expressionStatement3 &&
114+ expressionStatement3 . Expression is InvokeMethodExpression invokeMethodExpression &&
115+ invokeMethodExpression . MethodName == "AppendQueryDelimited" &&
116+ invokeMethodExpression . Arguments . Count == 5 &&
117+ invokeMethodExpression . Arguments [ 0 ] . ToDisplayString ( ) == "\" metadata\" " )
118+ {
119+ return new SingleLineCommentStatement ( "Plugin customization: remove unnecessary AppendQueryDelimited for metadata" ) ;
120+ }
121+ return statement ;
122+ } ) ;
123+
124+ // Rebuild the method body with the modified statements
125+ method . Update ( bodyStatements : statements ) ;
126+ }
127+
128+ return base . VisitMethod ( method ) ;
129+ }
130+ }
0 commit comments