Use GroqDoc and OpenRewrite to write your Javadoc
BitFlipper
It is important to have Javadoc comments on our Java interfaces so developers know how the interface was intended to be implemented. Lets use OpenRewrite and an LLM to add Javadoc comments to our interfaces!
GroqCloud
For this blog we have used the GroqCloud platform for its ease of use and low latency. You will need to generate an API key on their playground.
The simple HTTP client written for this recipe reads the API key from the GROQ_API_KEY environment variable. It uses the Mixtral 8x7b model as it has the largest context window of the default models on GroqCloud.
We also need a system message we can send to the LLM so it gives us a response in the desired output. Initially we got some issues with Mixtral including markdown formatting in the response. After some experimentation the following system message works well:
You have been hired as a Javadoc writer. The user will send you a java interface and you will write the Javadoc for the methods. Do not respond with anything other than a pure Javadoc string. Let me repeat, do not send anything other than a pure Javadoc string, including any and all markdown formatting.
Running GroqDoc Locally
The repository with the full source code used for this article can be found here https://github.com/Bit-Flipper/groq-doc. To run GroqqDoc for any maven project do the following:
Clone the groq-doc repo
git clonehttps://github.com/Bit-Flipper/groq-doc.gitcd groq-docInstall the recipe to your local maven repository
mvn install -DskipTestscdinto the maven project you want to document and run:
mvn -U org.openrewrite.maven:rewrite-maven-plugin:run \
-Drewrite.recipeArtifactCoordinates=dev.bitflippers:groq-doc:1.0-SNAPSHOT \
-Drewrite.activeRecipes=dev.bitflippers.groqdoc.GroqDoc # or GroqDocScanningRecipe
OpenRewrite Recipe Implementation
There will be two main stages we need to implement in the recipe. First, we need to extract the Java interface and send it to the LLM. Second, we need to parse the response from the LLM. Finally, we add the Javadoc to the LST.
Generating Javadoc Comments with Groq
In order for Mixtral to be able to generate Javadoc for a given Java interface it needs to have a textual representation of the interface. Each Tree node in the LST has a print method that returns the string representation of the node as it would appear in the Java source code. We call the print method on the J.CompilationUnit so we get the contents of the whole file and send the output to Groq.
public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext executionContext) {
// Only document interfaces
if (!Type.Interface.equals(classDecl.getKind())) {
return classDecl;
}
J.CompilationUnit compilationUnit = getCursor()
.dropParentUntil(p -> p instanceof J.CompilationUnit).getValue();
String interfaceString = compilationUnit.print(getCursor());
// Call Groq to generate Javadoc string
String generatedJavaDoc = generateJavadoc(interfaceString);
...
}
Parse Groq Response
Groq will respond with a version of the interface we sent but will now also include Javadoc comments e.g.
public interface Greetings {
/**
* @return a generic greeting to the world
*/
String helloWorld();
/**
* @return a generic farewell to the world
*/
String goodbyeWorld();
/**
* @return a personalised greeting to the name supplied
*/
String helloTo(String name);
/**
* @return a personalised farewell to the name supplied
*/
String goodbyeTo(String name);
}
The easiest way we can extract the Javadoc comments from the response is to get OpenRewrite to parse the interface and give us an LST. We can then traverse the LST and extract the Javadoc.DocComment nodes with a simple visitor.
private Map<J.MethodDeclaration, Javadoc.DocComment> extractGeneratedJavadocComments(ExecutionContext executionContext, String generatedJavaDoc) {
SourceFile lst = JavaParser.fromJavaVersion()
.build()
.parse(executionContext, generatedJavaDoc)
.toList()
.getFirst();
ExtractJavadocs javaDocVisitor = new ExtractJavadocs();
javaDocVisitor.visit(lst, executionContext);
return javaDocVisitor.JAVA_DOC_COMMENTS;
}
private static final class ExtractJavadocs extends JavaIsoVisitor<ExecutionContext> {
public final Map<J.MethodDeclaration, Javadoc.DocComment> JAVA_DOC_COMMENTS = new HashMap<>();
public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext executionContext) {
method = super.visitMethodDeclaration(method, executionContext);
if (method.getComments().isEmpty() || !(method.getComments().getFirst() instanceof Javadoc.DocComment)) {
return method;
}
JAVA_DOC_COMMENTS.put(method, (Javadoc.DocComment) method.getComments().getFirst());
return method;
}
}
Once we have extracted the Javadoc comments from the LLM response we need to put a reference to them in a convenient place for later. Placing a message on the cursor's message stack is ideal for this.
public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext executionContext) {
...
// Call Groq to generate Javadoc string
String generatedJavaDoc = generateJavadoc(interfaceString);
Map<J.MethodDeclaration, Javadoc.DocComment> docs = extractGeneratedJavadocComments(executionContext, generatedJavaDoc);
getCursor().putMessage(JAVA_DOCS_KEY, docs);
return super.visitClassDeclaration(classDecl, executionContext);
}
Modifying LST with Generated Comment Nodes
Once we've put the Javadoc.DocComment LST nodes on the message stack we need to retrieve the comment that matches this method. Unfortunately we can't do a simple map lookup as the equals methods won't match on the J.MethodDeclarations as they will have different UUIDs. Instead we will need to manually filter the map entries ourselves. Here the filtering is done by generating a string representation of the method signatures for comparison.
public J.MethodDeclaration visitMethodDeclaration(J.MethodDeclaration method, ExecutionContext executionContext) {
method = super.visitMethodDeclaration(method, executionContext);
// Only add Javadoc if there are no existing comments
if (!method.getComments().isEmpty()) {
return method;
}
// Get Javadoc LST node from message stack
Map<J.MethodDeclaration, Javadoc.DocComment> map = getCursor()
.getNearestMessage(JAVA_DOCS_KEY);
final J.MethodDeclaration finalMethod = method;
List<Comment> docComments = map.entrySet().stream()
// Must compare method signatures, .equals() won't work. The J.MethodDeclaration have different UUIDs
.filter(entrySet -> doMethodDeclarationsHaveSameSignature(entrySet.getKey(), finalMethod))
.map(entrySet -> (Comment) entrySet.getValue())
.toList();
return method.withComments(docComments);
}
Generated Javadoc Results
Do you know how the Apache Structs MVC framework works? No? Me neither. Sounds like we should read the documentation! Lets generate some.
See the generated Javadoc for the whole structs repository here: https://github.com/Bit-Flipper/struts/pull/2/files
Most of the Javadoc generated is quite good for example the Javadoc added to the Storage interface seems quite reasonable (coming from someone who does not know how structs works 😅).
There are some examples where the LLM does appear to be hallucinating though such as on the TagGenerator interface. Was the generate() part of structs since version 1.0? Also I'd like to meet the human named Your Name.
Will GroqDoc do Better with more Context?
Will Mixtral be able to generate better results with more context? Perhaps the LLM will be less likely to hallucinate if it has some examples of interface implementations. To get example implementations of the interfaces we can scan the rest of the structs repository for interface implementations. To do this we will need to convert the GroqDoc recipe to a scanning recipe.
OpenRewrite Scanning Recipe
See What is a ScanningRecipe? and How to Duplicate Classes with OpenRewrite articles for more information on how scanning recipes work.
The reason we will need to use a scanning recipe is so we can have the context of multiple source files while we are making changes to our LST representation of an interface. We will use the scanning phase and accumulator features of the scanning recipe to do this.
Lets change our recipe to use a map of fully qualified interface names to a list of example implementations like so class GroqDocScanningRecipe extends ScanningRecipe<Map<String, List<String>>>.
We will give the accumulator an empty map for its initial value.
public Map<String, List<String>> getInitialValue(ExecutionContext ctx) {
return new HashMap<>();
}
Next we can populate the accumulator map with all relevant example implementations as we find them. We do this by traversing every class declaration in the repository and store the source code of the classes if they implement an interface.
public TreeVisitor<?, ExecutionContext> getScanner(Map<String, List<String>> acc) {
return new JavaIsoVisitor<>() {
public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext executionContext) {
classDecl = super.visitClassDeclaration(classDecl, executionContext);
// If this class does not implement any interfaces we are not interested
if (classDecl.getImplements() == null || classDecl.getImplements().isEmpty()) {
return classDecl;
}
var implementsInterface = classDecl.getImplements().getFirst();
J.ClassDeclaration finalClassDecl = classDecl;
acc.compute(implementsInterface.getType().toString(), (k, v) -> {
if (v == null) {
v = new ArrayList<>();
}
// Add textual representation of implementation to accumulator
v.add(finalClassDecl.print(getCursor()));
return v;
});
return classDecl;
}
};
}
Finally we can pass the example implementations we found to Groq so the LLM will have more context to hopefully generate better Javadocs for us! During the editing phase we have access to the data stored in the accumulator.
public TreeVisitor<?, ExecutionContext> getVisitor(Map<String, List<String>> acc) {
return new JavaIsoVisitor<>() {
public J.ClassDeclaration visitClassDeclaration(J.ClassDeclaration classDecl, ExecutionContext executionContext) {
...
J.CompilationUnit compilationUnit = getCursor()
.dropParentUntil(p -> p instanceof J.CompilationUnit).getValue();
String interfaceString = compilationUnit.print(getCursor());
// Now we can access the accumulated example implementations
List<String> classDeclContext = acc.get(classDecl.getType().getFullyQualifiedName());
// ...and pass it to Groq
String generatedJavaDoc = generateJavadoc(interfaceString, classDeclContext);
Map<J.MethodDeclaration, Javadoc.DocComment> docs = extractGeneratedJavadocComments(executionContext, generatedJavaDoc);
...
}
};
}
Scanning Recipe Generated Javadoc
See the scanning recipe generated Javadoc for the whole structs repository here: https://github.com/Bit-Flipper/struts/pull/1/files
Does giving mixtral more context help generate better Javadoc with fewer hallucinations? In short, not really. We still see some high quality Javadoc generated like for the ClassFinderFactory.
Unfortunately we still see some hallucinations. With the MemberAccessValueStack interface we see the generated Javadoc talking about throwing NullPointerException or UnsupportedOperationExceptions which appears to be flat out wrong. The OgnlValueStack implementation of this interface does not throw either of these exceptions (see here).
Summary
Overall GroqCloud provides a really easy to use and performant API.
Getting an LLM to document your code can produce good results.
- Sometimes these results will have hallucinations in them so it is probably not practical to use in the real world (yet).
Giving the LLM more context in the form of example interface implementations does not appear to improve results or reduce hallucination frequency.
References
Support Our Work
Subscribe to my newsletter
Read articles from BitFlipper directly inside your inbox. Subscribe to the newsletter, and don't miss out.
Written by