Skip to main content

Annotations and LLMs

· 6 min read
Sanjeev Sarda
Sanjeev Sarda
High Performance Developer

Notes and ideas on annotations and LLMs. Using annotations in conjunction with LLM dev tooling as well as generating annotation processors with LLMs

Annotations

Annotations and LLMs​

What is an Annotation Processor?​

In Java we can use annotations to generate code and add functionality - this can be runtime or compile time.

@AutoService(Processor.class)
@SupportedSourceVersion(SourceVersion.RELEASE_21)
@SupportedAnnotationTypes("com.bhf.something")
public class SomethingProcessor extends AbstractProcessor {

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {

    }
}

There are a lot of subtleties to writing annotation processors. As part of the compilation process, there is an annotation processor step which calls the various processors you have configured in order.

An annotation processor can generate output which contains annotations, which in turn need to be processed. Processors are called in multiple rounds, which is why we have a RoundEnvironment passed to us in the process() which we have to implement for our own annotation processors.

We can implement the process method as follows:

annotations.forEach(annotation ->
    roundEnv.getElementsAnnotatedWith(annotation)
            .forEach(this::generateOutput));
    return true;

We return true to prevent subsequent processing of the particular annotation type we're dealing with in this processor. This is because a processor can handle multiple annotations via the @SupportedAnnotationTypes annotation.

Getting Fields from a Type level Annotation​

I wrote an annotation processor previously just to try and get a better idea how they worked - one thing I struggled with at the time was figuring out how to get the fields out, so this is as much an aside for myself as anything else:

void generateOutput(Element element) {
    var className = element.getSimpleName().toString();
    var packageName = element.getEnclosingElement().toString();

    var generatedType = "com.bhf.something";
    var outputName = className + generatedType;
    var outputFullName = packageName + "." + outputName;

    List<? extends Element> fields = element.getEnclosedElements()
            .stream().filter(e -> FIELD.equals(e.getKind())).toList();
}

Ways of Using Annotations and LLMs​

Domain Model and Service Boiler Generation​

Generating a polyglot domain model including things like UI components, unit and integration tests, Swagger spec etc.

Generating synthetic data examples for your domain model, either as part of the domain model generation, or an auto-generated service.

For example, in trading systems we may want to generate synthetic NoS, QuoteRequests or ExecReports which are tailored by asset class, client etc.

To generate the output in other languages you could use a prompt based approach (different prompts for different target languages and different underlying core models) or use an existing LLM to do the conversion from one language to another (convert Java to Typescript or Rust for example).

Developer Tooling​

Acting as a marker for other LLMs that you may already have plugged into your IDE.

A good example is Agni Annotations, which for full disclosure is written by me and still in the POC phase.

Part of that concept is that we "apply" prompts to subsets of the code using marker annotations which provide better context to the LLM, or enforce certain behavior like checks with the operator.

Prompt Development​

This would involve prompt development with preference based pairs and rankings e.g. for an annotation like @CPUConcerned we would provide preference pairs for code that is written in a good way if you're concerned about CPU usage vs code that is written in a bad way if you're concerned about CPU usage.

Output is further tuned using demonstration feedback from a human - we get a generated solution by the LLM and provide it with a demonstration of how we would have written it ourselves.

Custom LLMs​

A small custom model for this kind of task could be developed using a similar process to that of Microsoft's Phi series models, where the model is trained on sample code, a synthetic text book generated from a larger LLM, as well as a set of questions/coding exercises.

An LLM as an Annotation Processor​

The LLM is now the annotation processor and can also be designed to be a multi-step system.

You pass structured data to the LLM by using an annotation processor to get info about the annotated class.

You could also use an annotation processor to build a general schema sent to another service (using REST for example) - this service processes the schema using LLMs or whatever technique it wants.

This now decouples the language and technique of generating code from Java to any language you want - generate Rust from Python and have it put back in your generated-sources set in your Java workspace.

LLM Generated Annotation Processors​

This is using an LLM to actually write the annotation processor logic which will actually generate Java or another language using something like Java Poet or just straight up string concat.

You create and optimise prompts that will create annotation processors to output what you want.

A big advantage is that this could also reduce token consumption if you're using LLMs often for something that could be done by simpler or more traditional code generation techniques.

Considerations and Approach​

Prompt Engineering​

  • System prompt development process
  • Prompts to verify output vs. unit tests or actually compiling and testing (using an LLM tool function) vs. comparing bytecode generated
  • Guardrail prompts - ensuring safety of the annotation processor
  • Using a continuous vs discrete prompt optimisation approach

Prompt Enhancement​

  • Prompt evaluation - metrics via automation (run the code via tools), metrics via human rating, demonstration based feedback
  • Pair based optimization - preference pairs
  • Using multiple LLMs - low and high power LLMs, tools, as part of evaluating the output
  • Generating unit tests for the AP in advance
  • Evaluating correctness vs. optimality
  • Using MIPRO for evaluating the generation of annotation processor as these are workflows with multiple steps - an end to end evaluation and optimisation layer
  • Using embeddings to find similar annotation processors that we've generated
  • DICL for inference time optimisation

https://www.alibabacloud.com/blog/java-annotation-processing-tool---abstractprocessor_600185 - Java Annotation Processing Tool - AbstractProcessor, Alibaba Cloud Blog

https://www.youtube.com/watch?v=ja4is9oq37k - Annotation Processing in Java, Geekifik

https://arxiv.org/html/2405.11514v2 - Towards Translating Real-World Code with LLMs: A Study of Translating to Rust, Various, MPI-SWS, University of Bristol, Amazon Web Services, Inc

https://lokalise.com/blog/llm-code-translation/ - LLM Code Translation, Lokalise