Building AI-Powered Developer Tools at Jane Street

[on-hold music] My name is John Crepezzi, and I work on a team at Jane Street called AI Assistance.

Our group roughly, uh, is there at Jane Street to try to maximize the value that Jane Street can get from large language models. And I've spent my entire career, uh, in dev tools.

Before I worked at Jane Street, I was at GitHub for a long time, and then before that, I worked at a variety of other dev tools companies. And LLMs kind of present this really amazing opportunity in that they're so open-ended that we can build kind of anything that we can imagine.

And it seems like right now, the only thing moving faster than the progress of the models is kind of our creativity around how to employ them. Uh, at Jane Street, though, we've made some choices that make adoption of off-the-shelf tooling a little bit more difficult than it is for other companies.

And kind of the biggest reason that we have this problem is that we use OCaml as a development platform. For those not familiar with OCaml, it is a, a functional, very powerful language, but it's also incredibly obscure language.

Uh, it was built in France, and its most common applications are in things like theorem proving or formal verification. It's also used to write programming languages. Um, we use OCaml kind of for everything at Jane Street.

So just a couple quick examples. When we write web applications, of course, web applications have to be written in JavaScript. But instead, we write OCaml, and we use a library called Js_of_ocaml that is essentially a, uh, OCaml bytecode-to-JavaScript transpiler.

When we write plugins for Vim, those have to be written in Vim script. Uh, but we actually use a library called Vcaml, which again is a OCaml to Vim script transpiler.

And, uh, even people at the company that are working on FPGA code, they're not writing Verilog, they're writing in an OCaml library called Hardcaml.

Uh, so why are the tools on the market available not good for working with OCaml? I think it kind of comes down to a few primary reasons. The first and the most important is that models themselves are just not very good at OCaml.

And this isn't the fault of the AI labs. This is just kind of a by-product of the amount of data that exists for training. So it's-- there's a really good chance that the amount of OCaml code that we have inside of Jane Street is just more than, like, the total combined amount of OCaml code that there exists

in the world, uh, outside of our walls.

The second is that we've made things really hard on ourselves. Partially as a by-product of working in OCaml, we've had to build our own build systems. We built our own distributed build environment.

We even built our own code review system, which is called Iron.

We develop all of our software on a giant monorepo application, and just for fun, instead of storing that monorepo in Git, we store it in Mercurial. [laughing]

And, uh, s- at last count, sixty-seven percent of the firm uses Emacs instead of normal editors, maybe like VS Code. Uh, we do have people using VS Code, but Emacs is the most popular.

And the last thing is we're dreamers. I mean, kind of everyone in this room hopefully is, is a dreamer in a way. Uh, and what I mean by this is we want the ability to kind of take LLMs and apply them to different parts of our development flow and light up different parts.

So maybe we wanna use large language models to resolve merge conflicts or build better feature descriptions or figure out who reviewers for features be. And we don't wanna be hampered by the boundaries between different systems when we do that.

Over the next fifteen minutes, I'm gonna cover our approach to large language models at Jane Street, uh, particularly when it comes to developer tools. Um, I'm gonna talk about custom models that we're building and how we build them.

I'm gonna talk about editor integrations. So these are the integrations into, uh, to VS Code, Emacs, and Neovim. And I will talk about, uh, the ability that we've built over time to evaluate models and figure out how to make them perform best.

And I guess at first glance, it's not really obvious that training models at all is a good idea. I mean, it's a very expensive proposition. It takes a lot of time, and it can go wrong in a ton of different ways.

Who here has trained a model before or tried to train something like a model? Maybe you took a foundation model and trained on top of it.

Cool. We were more convinced after we read this paper. This is a paper from Meta about a project called CodeCompose. And in this paper, they detail the results fine-tuning a model specifically for use with Hack.

Uh, Hack is actually pretty similar to OCaml, uh, not in its, like, syntax or function, but really just in the fact that it's used primarily at one company and not really used much outside of that company, even though it's open source.

So... Oh, actually, a fun fact, Hack is implemented in OCaml. So I think that's just, like, a total coincidence, but... [laughing]

Uh, we were pretty naive early on. We read this paper, and we decided that it would be really cool if we could, uh, replicate the results. We thought we would just take a model off the shelf, we would show it a bunch of our code, and then we would get back a model that, uh, worked like the

original model but knew about our libraries and idioms. It turns out that's just not how it works. Uh, it's not that easy. In order to get good outcomes, you have to have the model see a bunch of examples that are in the shape of the type of question that you want to ask the model.

So we needed to first create a goal, a thing that we wanted the model to be able to do. And in our, in our world, the goal that we came up with was this: We wanted to be able to generate diffs given a prompt.

So what that means is we wanted ed-- a, a user inside of an editor to be able to write a description of what they want it to happen, and then have the model suggest a potentially multi-file diff.

So maybe you wanna modify the test file, an ML file, and an MLI, which is kind of like a header file.

We wanted the diffs to apply cleanly, and we wanted them to have a good likelihood of type checking after they had been applied. And we were kinda targeting this range of up to one hundred lines as, uh, an i-ideal zone of what we thought LLMs would actually be capable of.

And in order for that to work, we needed to collect data, like I was talking about before. We needed data of the training shape that looked just like the test shape.

And this is what that shape looks like for this task. You need to be able to collect a bunch of examples of what context the model would have had beforehand, and then some prompt of what you want the model to do, written hopefully in the same way that a human would write it, and then some diff that

would accomplish that goal. So context, prompt, diff, and we need a bunch of these examples. So how do we get these? How do we get these training examples?

Kind of the first place to look is Features. Features is, I mentioned a code review system that we built internally. This is what it looks like. It's called Iron.

Uh, Features are very similar to pull requests. I think you can just, you know, swap that term in your head. And Features at first glance have exactly the data they want.

On the Description tab, they have a human-written description of a change, and on the Diff tab, they have the code that accomplishes the goal of the developer.

But on closer look, they're not exactly what you want, right? The way that you write a feature description or a pull request description is really very different from what you might wanna say inside of an editor.

So you're not writing multiple paragraphs in the editor. You're just saying something like, "Fix that error that's happening right now," and that's just not how we write feature descriptions.

Another problem with these features or pull requests is that they're really large, right? Often it's, uh, a feature is five hundred lines or a thousand lines. So in order to use it as training data, we would need to have an automated way to pull features apart into individual smaller components that we could train on.

So we need smaller things than features. What are those? Well, maybe commits. Commits are smaller chunks than features. Uh, this is what a typical commit log looks like at Jane Street.

So this is not like a Git short log. This is literally just, like, an actual... I want you to look at this as an actual Git log. And where it says Summary Z, that's my commit message.

We don't really use commits the same way that the rest of the world uses them. So we use commits mostly as checkpoints between different parts, parts of a development cycle that you might wanna revert back to.

Commits don't have a description, and they also have the same problem in that they're not isolated changes. They're, they're a collection of changes. What we actually ended up with was a approach called workspace snapshotting.

And the way that that works is we take snapshots of developer workstations throughout the workday. So you can think like every twenty seconds, we just take a snapshot of what the developer's doing.

And as we take the snapshots, we also take snapshots of the build status. So as the build that's running on the box, we can see what the error is or whether the build is green.

And we can kinda notice these little patterns. If you have a green to red to green, that often corresponds to a place where a developer has made an isolated change, right?

You start writing some code, you break the build, and then you get it back to green, and that's how you make a change. Maybe this one, the red to green, this is a place where the developer encountered an error, whether that's a type error or a compilation error, and then they fixed it.

So if we capture the build error at the red state and then the diff from red to green, we can use that as training data to help the model be able to recover from mistakes.

The next thing we need is a description, and the way that we did that, we just used a large language model. So we had a large language model write a really detailed description of a change in, in as much words as it possibly could, and then we just kept filtering it down until it was something that was

around the right level of what a human would write.

So now we have this training data, and training data is kind of only half the picture of training a model. So you, you have the, the supervised training data, and then you need to do the second part, which is the reinforcement learning.

This is really where models get a lot of their power, right? We, we align the model's ability to what humans think is actually good code. So what is good code?

I guess on the surface, good code is, I mean, it's, it's code. It has to parse as code, meaning if a piece of code doesn't go through the OCaml parser and come out with a green status, that is, that is not good code [chuckles] I would say by most definitions.

Uh, good code in OCaml, because it's statically typed, is code that type checks. So we wanna have good code be code that when it is applied on top of a, a base revision, can go through the type checker, and the type checker agrees that the code is valid.

And of course, the, the gold standard is that good code is code that compiles and passes tests. So ideally, in, during the reinforcement learning phase of a model, you could give the model a bunch of tasks that are, like, verifiable.

We have the model perform some, some edit, and then we check whether or not it actually passes the test when applied to the code.

So we did that. Uh, we've done this as part of our, our training cycle, and we built this thing that is called, uh, CES. It's the Code Evaluation Service.

You can think of it kinda like a build service, except with a slight modification to make it much faster, and that's that first we pre-warm a build. It sits at a, a revision and is green, and then we have these workers that all day just take diffs from the model, they apply them, and then we determine whether

the build status turns red or green, and then we report that error or, or success back up to the build function. And through continued use of this service over the course of, like, months, we're able to better align the model to write code that actually does compile and pass tests.

It turns out this exact same setup is the one that you would want for evaluation. So if you just hold out some of the RL data, you can also use it to evaluate model's ability to write code.

Kinda looks like this. You give the model a problem, you let it write some code, and then you evaluate whether or not the code that it writes actually works.

And training is hard, and it can have kind of, uh, catastrophic but hilarious results. So at, at one point, we were training a code review model, and this is a totally separate model, but the idea was we wanna be able to give some code to this model and have it do a first pass of code review just

like a human would do to try to save some of the toil of, of code review. We trained this model, we put a bunch of data into it, we worked on it for months, we were real excited, and we put our first code in for, uh, for code review through the automated agent.

It spun for a bit, and it came back with something along the lines of, "Mm, I'll do it tomorrow."

And [laughs] like, of course, it did that because it's trained on a bunch of human examples, and humans write things like, "I'll do things tomo- or I'll do this tomorrow."

Uh, so it's, it's, you know, not very surprising. So having evaluations that are meaningful is kind of a cornerstone of making sure that models don't go off the rails like this, and you don't waste a bunch of your time and money.

In the end though, the real test of models is whether or not they work for humans. So I'm gonna talk a little bit about the editor integrations that we've built to expose these models to developers at Jane Street.

Kind of when we were starting building these integrations, we had three ideas in mind. The first idea was, wow, we support three editors. We have Neovim, VS Code, and Emacs, and we really don't wanna write the same thing three times.

So ideally, we don't wanna write all the same context building strategies and all of the same prompting strategies. We wanna just write it once. The second is that we wanted to maintain flexibility, so we had a model that we were using at the time, uh, that was not a fine-tuned model.

We were pretty convinced that a fine-tuned model was in our future. We wanted the ability to do things like swap the model or swap the prompting strategy out. And last, we wanted to be able to collect metrics.

So in a developer, uh, in their, in their editor, developers care about latency, and they care about whether or not the diffs actually apply. So we wanted to get kind of on-the-ground, real experience of whether or not the diffs really were meaningful for people.

This is the simplified version of the architecture that we settled on for this service, the AI development environment. Essentially, you have LLMs on one side, and then Aid handles all of the, uh, ability to construct prompts and to construct context and to see the build status, and then we are able to just write these really thin lo-

layers on top of Aid, uh, for each of the individual editors. And what's really neat about this is that Aid sits as a sidecar application on the developer's machine, which means that we w- when, when we wanna make changes to Aid, we don't have to make changes to the individual editors and hope that people restart their editors.

We can just restart the Aid service on all of the boxes. So we restart Aid, and then everyone gets the most recent copy.

Uh, this is an example of Aid working inside of VS Code. So this is the sidebar in VS Code, very similar to something like Copilot, except this thing allows you to, uh, ask for it and get back multi-file diffs.

Uh, and you can see it kinda looks like what you'd expect in VS Code. It's, it's, you know, a visual interface that lays things out really nicely.

This is what we built in Emacs though. So in Emacs, developers are used to working in text buffers. They move around files. They wanna be able to copy things the normal way that they copy things.

So we actually built the Aid experience in Emacs into a markdown buffer. So users can move around inside this markdown buffer, they can ask questions, and then there are key binds that essentially append extra content to the bottom of the markdown buffer.

Aid's architecture lets us plug various pieces in and out, like I mentioned. Uh, so we can swap in new models. We can, uh, make changes to the context building.

We can add support for new editors, which I think probably sounds far-fetched, but we're-- this is something we're actually just doing right now.

Uh, and we can even add domain-specific tools. So different areas of the company can supply tools that are available inside of the editors, and they kind of end up in all the editors without hav- having to write individual integrations.

Aid also allows us to A/B test different approaches, so we can do something like send fifty percent of the company to one model and fifty percent to another and then determine which one gets the higher acceptance rate.

Aid is kind of a, an investment that pays off over time. Every time something changes in large language models, we're able to change it in one place downstream of the editors and then have it a-available everywhere.

And things change, like, really often, and we need to be ready, uh, when things change. What I, what I've had time to show you today is only a small portion of what my team is doing.

Uh, we've got a lot of other things going on, so we're finding new ways to apply RAG inside of the editors. We're applying similar approaches to what you've seen here to large scale, uh, multi-agent workflows.

We are working with reasoning models more and more. But the approach is the same through all of these. We keep things pluggable. We lay a strong foundation to build on top of, and we build the ways for the rest of the company to add to our experience by adding more domain-specific tooling on top of it.

Uh, if you think what I've said is interesting and you wanna talk more about this, I would love to hear from you. You can just find me outside, and thank you for your time. [upbeat music]