feat: method-by-method OAS doc generation strategy #6023
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,23 +5,115 @@ | |
| * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause | ||
| */ | ||
|
|
||
| import * as fs from 'fs'; | ||
| import { DocumentSymbol } from 'vscode'; | ||
| import * as yaml from 'yaml'; | ||
| import { | ||
| ApexClassOASEligibleResponse, | ||
| ApexClassOASGatherContextResponse, | ||
| ApexOASClassDetail, | ||
| ApexOASMethodDetail, | ||
| PromptGenerationResult, | ||
| PromptGenerationStrategyBid, | ||
| OpenAPIDoc | ||
| } from '../../openApiUtilities/schemas'; | ||
| import { IMPOSED_FACTOR, PROMPT_TOKEN_MAX_LIMIT, RESPONSE_TOKEN_MAX_LIMIT, SUM_TOKEN_MAX_LIMIT } from '.'; | ||
| import { GenerationStrategy } from './generationStrategy'; | ||
| import { prompts } from './prompts'; | ||
| export const METHOD_BY_METHOD_STRATEGY_NAME = 'MethodByMethod'; | ||
| export class MethodByMethodStrategy extends GenerationStrategy { | ||
| async callLLMWithPrompts(): Promise<string[]> { | ||
| try { | ||
| const llmService = await this.getLLMServiceInterface(); | ||
|
|
||
| // Filter valid prompts and map them to promises | ||
| const responsePromises = this.prompts.filter(p => p?.length > 0).map(prompt => llmService.callLLM(prompt)); | ||
|
|
||
| // Execute all LLM calls in parallel and store responses | ||
| await Promise.allSettled(responsePromises).then(results => { | ||
|
There was a problem hiding this comment. use Promise.allSettled to populate the methods as many as possible and not be blocked by failed requests |
||
| results.forEach((result, index) => { | ||
| if (result.status === 'fulfilled') { | ||
| this.llmResponses.push(result.value); | ||
| } else if (result.status === 'rejected') { | ||
| console.log(`Promise ${index} rejected with reason:`, result.reason); | ||
| } | ||
| }); | ||
| }); | ||
|
|
||
| return this.llmResponses; | ||
| } catch (error) { | ||
| const errorMessage = error instanceof Error ? error.message : String(error); | ||
| throw new Error(errorMessage); | ||
| } | ||
| } | ||
|
|
||
| public async generateOAS(): Promise<string> { | ||
| const oas = await this.callLLMWithPrompts(); | ||
| if (oas.length > 0) { | ||
| try { | ||
| const combinedText = this.combineYamlByMethod(oas); | ||
| return combinedText; | ||
| } catch (e) { | ||
| throw new Error(`Failed to parse OAS: ${e}`); | ||
| } | ||
| } else { | ||
| throw new Error('No OAS generated'); | ||
| } | ||
| } | ||
|
|
||
| cleanYamlString(input: string): string { | ||
|
There was a problem hiding this comment. consider adding the There was a problem hiding this comment. this is resolved by Cristi's tip above |
||
| return input | ||
| .replace(/^```yaml\n/, '') // Remove leading triple backtick (if any) | ||
| .replace(/\n```$/, '') // Remove trailing triple backtick (if any) | ||
| .replace(/```\n\s*$/, '') // Remove trailing triple backtick with new line (if any) | ||
| .trim(); // Ensure no extra spaces | ||
| } | ||
|
|
||
| combineYamlByMethod(docs: string[]) { | ||
| const combined: OpenAPIDoc = { | ||
| openapi: '3.0.0', | ||
| info: { | ||
| title: this.context.classDetail.name, | ||
| version: '1.0.0', | ||
| description: `This is auto-generated OpenAPI v3 spec for ${this.context.classDetail.name}.` | ||
|
There was a problem hiding this comment. the description is templated before we elaborate how it should look like |
||
| }, | ||
| paths: {}, | ||
| components: { schemas: {} } | ||
| }; | ||
|
|
||
| for (const doc of docs) { | ||
| const yamlCleanDoc = this.cleanYamlString(doc); | ||
| let parsed = null; | ||
| try { | ||
| parsed = yaml.parse(yamlCleanDoc) as OpenAPIDoc; | ||
| } catch (error) { | ||
| throw new Error(`Failed to parse YAML: ${error}`); | ||
| } | ||
| // Merge paths | ||
| for (const [path, methods] of Object.entries(parsed.paths)) { | ||
| if (!combined.paths[path]) { | ||
| combined.paths[path] = {}; | ||
| } | ||
| Object.assign(combined.paths[path], methods); | ||
| // explicitly define openrationId if missing | ||
| for (const [method, props] of Object.entries(combined.paths[path] as object)) { | ||
| if (props?.operationId === undefined) { | ||
|
There was a problem hiding this comment. I tried several different prompts to remind LLM always enabling operationId, but found out it still failed to fulfill my will constantly. So I added this mandatory check. There was a problem hiding this comment. an example how it could fail @RestResource(urlMapping='/simpleResource/*')
global with sharing class SimpleRestResource {
@HttpGet
global static String doGet() {
return 'Hello World';
}
@HttpPost
global static String doPost(String message) {
return 'You posted: ' + message;
}
@HttpPut
global static String doPut(String message) {
return 'You put: ' + message;
}
@HttpDelete
global static String doDelete() {
return 'Goodbye';
}
}
|
||
| combined.paths[path][method].operationId = path.split('/').pop() + '_' + method; | ||
| } | ||
| } | ||
| } | ||
| // Merge components | ||
| if (parsed.components?.schemas) { | ||
|
There was a problem hiding this comment. @mingxuanzhangsfdx should this check if the schema already exists, as is being done above with paths |
||
| for (const [schema, definition] of Object.entries(parsed.components.schemas)) { | ||
| if (!combined.components!.schemas![schema]) { | ||
| combined.components!.schemas![schema] = definition as Record<string, any>; | ||
| } | ||
| } | ||
| } | ||
| } | ||
| return yaml.stringify(combined); | ||
| } | ||
|
|
||
| llmResponses: string[]; | ||
| public async callLLMWithGivenPrompts(): Promise<string[]> { | ||
| let documentContent = ''; | ||
|
|
@@ -45,6 +137,10 @@ export class MethodByMethodStrategy extends GenerationStrategy { | |
| callCounts: number; | ||
| maxBudget: number; | ||
| methodsList: string[]; | ||
| methodsDocSymbolMap: Map<string, DocumentSymbol>; | ||
| methodsContextMap: Map<string, ApexOASMethodDetail>; | ||
| documentText: string; | ||
| classPrompt: string; // The prompt for the entire class | ||
|
|
||
| public constructor(metadata: ApexClassOASEligibleResponse, context: ApexClassOASGatherContextResponse) { | ||
| super(); | ||
|
|
@@ -56,6 +152,19 @@ export class MethodByMethodStrategy extends GenerationStrategy { | |
| this.maxBudget = SUM_TOKEN_MAX_LIMIT * IMPOSED_FACTOR; | ||
| this.methodsList = []; | ||
| this.llmResponses = []; | ||
| this.methodsDocSymbolMap = new Map(); | ||
| this.methodsContextMap = new Map(); | ||
| this.documentText = fs.readFileSync(new URL(this.metadata.resourceUri.toString()), 'utf8'); | ||
| this.classPrompt = this.buildClassPrompt(this.context.classDetail); | ||
| } | ||
|
|
||
| buildClassPrompt(classDetail: ApexOASClassDetail): string { | ||
| let prompt = ''; | ||
| prompt += `The class name of the given method is ${classDetail.name}.\n`; | ||
| if (classDetail.comment !== undefined) { | ||
| prompt += `The comment of the class is ${classDetail.comment}.\n`; | ||
| } | ||
| return prompt; | ||
| } | ||
|
|
||
| public bid(): PromptGenerationStrategyBid { | ||
|
|
@@ -65,10 +174,15 @@ export class MethodByMethodStrategy extends GenerationStrategy { | |
| }; | ||
| } | ||
| public generate(): PromptGenerationResult { | ||
| const list = (this.metadata.symbols ?? []).filter(s => s.isApexOasEligible); | ||
| for (const symbol of list) { | ||
| const methodName = symbol.docSymbol.name; | ||
| this.methodsDocSymbolMap.set(methodName, symbol.docSymbol); | ||
| const methodDetail = this.context.methods.find(m => m.name === methodName); | ||
| if (methodDetail) { | ||
| this.methodsContextMap.set(methodName, methodDetail); | ||
| } | ||
|
|
||
| const input = this.generatePromptForMethod(methodName); | ||
| const tokenCount = this.getPromptTokenCount(input); | ||
| if (tokenCount <= PROMPT_TOKEN_MAX_LIMIT * IMPOSED_FACTOR) { | ||
|
|
@@ -96,6 +210,44 @@ export class MethodByMethodStrategy extends GenerationStrategy { | |
| } | ||
|
|
||
| generatePromptForMethod(methodName: string): string { | ||
| let input = ''; | ||
| const methodContext = this.methodsContextMap.get(methodName); | ||
| input += `${prompts.SYSTEM_TAG}\n${prompts.METHOD_BY_METHOD.systemPrompt}\n${prompts.END_OF_PROMPT_TAG}\n`; | ||
| input += `${prompts.USER_TAG}\n${prompts.METHOD_BY_METHOD.USER_PROMPT}\n`; | ||
| input += '\nThis is the Apex method the OpenAPI v3 specification should be generated for:\n```\n'; | ||
| input += this.getMethodImplementation(methodName, this.documentText); | ||
| input += `The method name is ${methodName}.\n`; | ||
| if (methodContext?.returnType !== undefined) { | ||
| input += `The return type of the method is ${methodContext.returnType}.\n`; | ||
| } | ||
| if (methodContext?.parameterTypes?.length ?? 0 > 0) { | ||
| input += `The parameter types of the method are ${methodContext!.parameterTypes.join(', ')}.\n`; | ||
| } | ||
| if (methodContext?.modifiers?.length ?? 0 > 0) { | ||
| input += `The modifiers of the method are ${methodContext!.modifiers.join(', ')}.\n`; | ||
| } | ||
| if (methodContext?.annotations?.length ?? 0 > 0) { | ||
| input += `The annotations of the method are ${methodContext!.annotations.join(', ')}.\n`; | ||
| } | ||
| if (methodContext?.comment !== undefined) { | ||
| input += `The comment of the method is ${methodContext!.comment}.\n`; | ||
| } | ||
| input += this.classPrompt; | ||
| input += `\n\`\`\`\n${prompts.END_OF_PROMPT_TAG}\n${prompts.ASSISTANT_TAG}\n`; | ||
|
|
||
| return input; | ||
| } | ||
|
|
||
| getMethodImplementation(methodName: string, doc: string): string { | ||
| const methodSymbol = this.methodsDocSymbolMap.get(methodName); | ||
| if (methodSymbol) { | ||
| const startLine = methodSymbol.range.start.line; | ||
| const endLine = methodSymbol.range.end.line; | ||
| const lines = doc.split('\n').map(line => line.trim()); | ||
| const method = lines.slice(startLine - 1, endLine + 1).join('\n'); | ||
| return method; | ||
| } else { | ||
| throw new Error(`Method ${methodName} not found in document symbols`); | ||
| } | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@mingxuanzhangsfdx what happens if there is a tie on number of calls?