Making LLMs Do What You Want
Fri Oct 25, 2024Sorry I've been gone a while. Startup life, yo.
Audibly Clunky Segue
So a use case I have fairly commonly is that I want a model to actually do what I want. Not like, "approximately what I want" or "close enough to what I want that it doesn't raise exceptions", or "do what I want impressionistically enough that I have to call the function String -> String
ala markdown
transformation".
I mean like, I want to write a prompt, provide some data and be sure that the thing I get out is Good Enough to do Real Work with. I'll admit that this isn't always possible. There are situations where you have to do the String -> String
thing, or that are complicated enough that you'll have to tolerate some error rate above 0%.
But there are some "special cases" where it's kinda hard to write a solution to the problem yourself, but it's either trivial or really easy to check whether a solution someone hands you works to within the tolerance you care about. Hello, everyone who's been in the NP space.
Pseudo
Suppose you had a function
generate :: SystemPrompt -> Prompt -> Response
All three of those arguments are effectively aliases for strings.
As an aside, in case you didn't know, this is the whole story. LLMs are systems where you pour strings in, mix them around with enough GPUs and get strings out. If you don't believe me, go read the docs on how Claude tool use actually works "under the hood" in heavy finger-quotes.
Sorry, suppose you had that generate
function. You could trivially write
generateChecked :: (String -> Maybe a) -> SystemPrompt -> Prompt -> ?Int -> Maybe a
generateChecked (transformFn, system, prompt, retries=5) =
let res = generate system prompt
transformed, success = transformFn res
in if transformed:
Just transformed
else
retry upTo retries or Nothing
which is going to take the output, validate it against the transformer you've provided, and retry until it either gets to a valid value or retries too much. If you read that type signature, it's pretty obvious.
A more specific use case here is something I frequently want to ask an LLM: "Hey, evaluate this bunch of weirdly structured, natural language data and return a value with the following type signature so I can incorporate your output into some code I'm writing". The simplest here might be
generateJson :: SystemPrompt -> Prompt -> ?Int -> Maybe JSONValue
generateJson (system, prompt, retries=5) =
generateChecked maybeParseJson system prompt retries
There. Now you can call your LLM, parse the response it gave you into JSON. If it fails, that is if the response isn't properly formatted JSON, try again some specified number of times. Another use case might be
generatePageSelector :: Page -> Prompt -> ?Int -> Maybe String
generatePageSelector (page, prompt, retries=5) =
generateChecked #(querySelector page %) prompt retries
And again. Now you can ask your LLM to look at a web page, pick out an element by CSS selector, ensure that such a selector actually exists in the given page, and return it if it does. And retry some number of times if the first try doesn't hit. The code that relies on that response can rely that it's either looking at a selector that actually exists, or that it very clearly and explicitly failed on retries.Python
What you thought this was just woolgathering?
Because we're dealing with a weakly typed system, we can't rely on Maybe. We have to go the Lispy-route of returning multiple values in part of this.
def loadch(resp):
try:
return (
json.loads(
(resp.strip().removeprefix("```json").removesuffix("```").strip())
),
True,
)
except (TypeError, json.decoder.JSONDecodeError):
pass
return None, False
We're using the standard json.loads
, but stripping some common prefix/suffix strings that come back even when you try to force most LLMs to really just return readable JSON. We have to return a secondary value. Because None
and False
are both valid python values and also valid JSON values, and Python doesn't have a native equivalent of Maybe
, we'll have to work with what we've got here. So we return a second Boolean
value to specify whether we successfully parsed. This lets us disambiguate
None, False # we failed to parse all of the retry expressions we got
from
None, True # we successfully parsed the expression, and its value was NULL, which translates to python None
Ok, so we've got a way to parse generated strings into JSON. Lets get a way of getting generated text to begin with. This'll be an example using ollama
, but the repo will implement compatibility for all the big ones. At minimum, I'm planning to support Claude, Gemini and ChatGPT.
class Ollama(LLMMixin):
def __init__(self, model, ollama_server):
self.model = model
self.server = ollama_server.rstrip("/")
def generate(self, system, prompt):
res = requests.post(
f"{self.server}/api/generate",
json={
"model": self.model,
"stream": False,
"prompt": f"SYSTEM PROMPT: {system} PROMPT: {prompt}",
},
)
if res.status_code == 200:
return LLMResult(res, 200, res.json()["response"].strip())
return LLMResult(res, res.status_code, None)
The LLMResult
type is just a namedtuple
I've set up so that you can do dot notation on it, but realistically, it could be a dictionary. Just ignore it. The generate
function there is obvious. But I snuck in LLMMixin
. Lets take a look at that:
class LLMMixin:
def generate_checked(self, transformFn, system, prompt, retries=5):
for i in range(retries):
res = self.generate(system, prompt)
transformed, success = transformFn(res.content)
if success:
return LLMResult(res.raw, res.status_code, transformed)
return LLMResult(res.raw, 500, None)
def generate_json(self, system, prompt, retries=5):
return self.generate_checked(loadch, system, prompt, retries=retries)
That generate_checked
is what we've been building to. Again, python
, so we don't have the cool notional type system that I think in most of the time, and that makes us deal with the return tuple from the transform function. But you should see what the mapping between this and the earlier pseudocode is.
...
def generate_checked(self, transformFn, system, prompt, retries=5):
for i in range(retries):
res = self.generate(system, prompt)
transformed, success = transformFn(res.content)
if success:
return LLMResult(res.raw, res.status_code, transformed)
return LLMResult(res.raw, 500, None)
...
So, retries
times, call generate
, call the given transformFn
on the result, if the transformation succeeds short circuit and return the transformed code. If you fail, return the raw
input and None
. As an aside, you can see why I want you to ignore LLMResult
for the moment; I don't think I'm keeping that status_code
field long term.Goals
The ultimate goal here is better, more composeable scaffolding. I think I've tipped my hand a bit in terms of what I'm currently using this for by showing you the concrete JSON and CSS selector examples. Details left as an exercise for the reader, at least for now. I've got a few bigger ones in mind though. Keeping cards close to my chest for now, but you can probably guess at least one or two of them based on my github.
As always, I'll let you know how it goes.