Self-correcting code assistants with Codestral

Key Takeaways

hey this is Lance from Lang chain so mraw released code STW today which is a code generation model um which I'm actually really excited about so it's really good at code generation tasks like fill INE middle or code completion it's trained on a programming language it has an instruct version that supports tool use but one of the reasons why I really like code generation models and I've actually done quite a bit of work with them is that they're just very generally useful so lots of companies for example want C custom code assistance that might combine like some documentation Plus Code gener ation um at Lang chain for example we have something called chat Lang chain it's basically QA over our docs it can produce functioning code blocks for users based on questions um and one of the other things is cool about code is really easy to evaluate it's really easy to test does this code actually execute or not um and so a really powerful idea related to code generation was put out a few months ago um from the folks at codium Ai and carboi summarized it really nicely here in that this idea of flow engineering for code for code generation is really powerful and the idea shown in the paper the alpha codium work and kind of highlighted here in this tweet from karpathy is simply that if you produce a code solution you can really easily check it in line kind of as mentioned here it's pretty easy to evaluate code at the minimum does it execute do the Imports work um in the maximum case do you have like a an actual solution do unit test but in any case the point is code is very easy to test and you can actually test it in your inference flow so you produce a generation you then test the code if it fails you can loop back and try again and this idea of kind of a a code generation flow was shown in the paper to produce much better results and it's something that I want to show today uh using Cod strw uh in a really simple test case so this is something I've done a little bit in the past and I found to be extremely effective it's a very simple idea but here's the basic flow that we will kind of highlight so I want to be able to take a question related to code generation pass it to the model so pass it to coal and have cool produce a solution um now what I'm going to do is I'm going to use a function calling or tool use with codol to produce an output object that has three things a preamble stating like here is the problem I'm trying to solve the Imports and the code itself and what I'm going to do is I'm going to show how it's really easy to incorporate some simple code checks like do the Imports work the code execute if either fail like there's a bug in the code then I'm going to show how to loot back and retry and this simple kind of like check retry Loop is a way to significantly improve the the accuracy and and kind of usability of code generation models I'm going to show how to do that right now so to kick this off I have a notebook here I've done a few pip installs I've just set my mistol API key that's really it and I'm also going to use Lang Smith or tracing of course this is optional I'm going to set envir IR ment variable for my line chain project which will basically will indicate where all my traces will go and this is just like the the kind of flow we want to lay out so we basically want to use cod strol to take a user question produce a solution and we want to test that solution if it passes our test return to the user if it doesn't try again that's all we want to do so what I'm going to do here is first let me just show some very basic components so first let's just talk about how to actually use code here's my general prompt I basically just I'm going to tell the model your code assistant ensure that all the code can be executed with all the Imports and variables defined structure your answer in three ways give me a preamble or a prefix describing the code solution give me the Imports give me a function and code block so I'm going to ask for those three things now here's where tool use comes in I can actually Define the schema of the output that I actually want and what I can do is I can bind that using a lang chain very convenient with structured output I can basically bind that to the llm and then this chain will invoke the LM using the structured output now here's how that actually works under the hood basically this object that passes a pantic object is converted into function schema form a STW and it's then passed or bound to the llm so the llm has access to this function and it knows the schema that it should return when that function call or tool is invoked so basically what happens is I can take a user question the function is invoked and then the llm knows to produce an output that adheres to my schema and this will basically be a Json string again remember llm is just string to string so it's going to be a Json string and then under the hood with this with structured output thing that I'm using from Lang chain we apply an output parser that basically pantic parti should take a Json string convert it back to pantic object so that's it that's all that's going on um but I'll show you how this is really cool so I'm defining this object this is what I want to get out here's my chain now let's test this out write a function for Fibonacci um I passed it in as a user question um so that's it now this is running great and we see a result now here's what's cool if you look at this result object it actually is a code object so it basically it's pantic object following the scheme we specify here it has a prefix boom it has some imports actually in this case none and then it has the code block that's it so we'll see why this is really useful in a little bit but want to introduce that idea of basically we can use cod strol with tool use to produce structured outputs which is generally very useful and in the particular for this notion of kind of like inline self-correction is extremely useful cool so that's that first piece now what I'm going to introduce here is Lang graph so Lang graph is a library from the Lang chain team and we've used this a number of other videos um and I've used this kind of extensively in general uh to build flow flows and this is an example of a flow the main characteristic of the flow that I highlight here that Lang graph is really well suited for is anything with a cycle so anything with feedback basically what it's saying is I want every time I run my app I want to do this code generation produce a structured output do some kind of code checks make a decision based on the outcome of those code checks feedback if they fail finish if they pass that you can think of as like a very simple kind of like workflow um and L graph is a great way to build these kinds of workflows and we'll see why so the first thing I need to specify with L graph is just simply the graph State now this is just a thing that lives throughout the lifetime of My Graph it basically represents all information that's shared across what you might call these nodes so in this case I have two particular nodes and you might call this an edge so this is kind of where I'm making a decision um so State lives across these nodes and edges so that's really it so I'm going to Define my state it's going to attain some information that's relevant to the flow I just talked about so it's contain an error message it's going to contain my final generation it's going to contain the messages that are being passed to my llm and this will all become a little bit more clear as we go forward so here I'm going to lay out this is basically the nodes and the edges of My Graph now what you'll see is for every node here's my generate node and that's what we laid out here generate um it's going to take in the state and the nodes just modify the state in some way so that's how to think about the nodes so in this case I take in the state I unpack the state into like some messages uh some of iterations an error message these are things we're going to use throughout our graph um so then what I'm going to Simply do is compute a code solution so I'm going to look at my messages in and I'm going to generate a solution now remember that's exactly what we did up here so this is actually nothing new remember look at this this is just we Define a set of messages invoke our code gen chain get an output same thing we're going to be doing in our graph so this is nothing exotic we've actually already tested this and once that runs I'm just going to pend that output of code solution to my messages okay so you know again here's my attempt to solve the problem I'm just going to take the codes the prefix the Imports and the code I'm just going to add that as a new message I'm going to increment my iterations we'll use as iterations to determine when to stop um and I'm just going to return then my state with a few things here first is going to be my code solution My Generation that's it then it's going to be my my stack of messages which is basically just pended to and then the number of iterations that's really it that's it so nice and easy there now the code check is the second kind of big node that we're going to be working with so our first node is Generation the second node is our code checks we just saw Generation generation can return the the generation with the three pieces the Preamble The Imports and the code block now it's going to be passed to code check so code check is really anything we want to be we can do any kind of checks on this code now maybe in the best case with some kind of unit test we could run I'm going to show you the simplest possible code check that we might want to do so in this particular particular case what I'm going to do is I'm I'm going to get the code solution from our state remember we wrote that out to state so the generation contains our code solution and in this node I just pick it back up from State you know State's pass every node I get the code solution I EX exract the three pieces and we just showed that above so I get the prefix the Imports and the code and now all I'm going to do is simply just test execution do does Imports execute if not I'm going to throw a flag or I'm going to kind of flow throw a message here code import failed I'm going to take an error message I'm going to pen that error message to our messages object and I'm going to return that that's it and alternatively if that passes I'm going to go ahead and try the whole thing so I'm going to combine the Imports and the code um I'm going to go ahead and execute the the code and again if that fails I'm going to basically return another error message um and now if there's no errors then that's great I confirm that you know there's no test failures I've set this error flag to no and everything else the same as before I return the messages I return iterations I return code Solution that's it now this is the final bit all we're going to do here is decide whether or not to finish this is basically our little conditional Edge which we talked about here and all this needs to to do because we wrote that error flag to state so again remember we wrote error no if none of these tests passed we wrote error yes if either one does right there and if that's the case all we need to do then is get our get our error from the state um if it's no or we've exceed the Max iterations uh then we just go ahead and finish um and if yes then we go back to generate that's it and that's really it so we're just going to find all of those pieces and we're basically almost done here let's just build this graph now this is how in L graph you can actually assemble your workflow all I need to do is take that function I defined generate add it as a node take the function we defined code check add it as a node again this is my like this is the state graph um and I just build the graph here set my entry point as generate um add an edge and then uh so basically I go from generate to check code and I go from check code to basically I decide to finish based upon this logic right here and basically if if it returns end then I end if it returns generate I go back to generate okay um so that's really the Crux of all you need to do and I can go ahead and run that and actually this will draw my graph for me using this little display feature right here so we can see we start we go to generate we go to code check optionally depending on what happens from the code check based onal Edge will'll go back to try to regenerate um or if there's no errors we go to end so that's really it it's pretty nice um and the one thing I'll just make a note of is as we're going through this flow we're actually appending to our messages and so basically if the if there's an error we're appending that failure to our messages and we're basically telling llm here is the failure reflect on this error um State what you think went wrong and try again so that's really it nice and easy and there's our flow cool now let's try this out here's like the simplest possible um you know kind of problem write a Python program that prints hello world right so let's try this out I'm going to run this and what's kind of nice is I have this kind of this this nice kind of formatting stuff you can kind of see the input um yeah okay so R program that prints hello world generating code solution and then here's here's kind of my attempt to solve it um here's actually the the Imports none the code here's the code it goes through the checks no test failures and it ends cool now what it's nice is I can go over to lsmith and I have this project right here now this project actually lays out exactly what we just did but we can actually dig into each piece so here we went we start our graph we went to generate I go to again this is using Cod strol model so here's here is my U you know human message or my question in um this is showing that it does indeed invoke my function so that's great um here's the prefix here's the Imports here's the code block um it uses a pantic tools parser to basically write that out as a pantic object we talked about that previously um and then here's the code check so basically it goes through the the various code checks and you can kind of check all these here um and then it goes through the decision to finish um and in this particular case because none of the code checks failed and finished so that's great so this is a good example of like of kind of how the flow Works in a very simple test case now let's try something this a little more sophisticated cool so in this case I'm basically asked to vectorize a function I give it a function um I show me ask it to show me a test case with this with this actually working okay so we can see that it kicks off the flow here I want to vectorize a function um here's the inmt to solve the problem uh so here's kind of the initial solution now what we see here is your solution failed the code execution test it did not Define image reflect on the error attempt to solve it here's my attempt to solve a problem the error C because variable the the variable image is not defined to solve this problem so you make it kind of reflect on its error and try again so it goes and tries again we see it fails again for a different reason isolution failed the C secution test it could not broadcast cast um uh 505050 into a shape 50-50 53 okay um so here's my here's the attempt to solve the problem it kind of explains itself it goes back through and then all the code tests now pass um so that's basically it and then finish so this showcases how you can use code generation using new code STW model stol with self correction using Lang graph and what we showed in in general is the ability to to perform code generation and perform arbitrary checks on the output of the generation itself if any checks fail Loop them back use a message cue to accumulate over time uh or over iterations in this flow the various errors and then pass them back to the LM to attempt to self-correct and we seen we've seen this work pretty effectively in a very simple test case but I've actually seen this work really well in the case of uh code generation with Rag and and uh we have some other uh resources on that which I'll share later thanks