Simple Hack to always generate JSON from LLM
Trick to always get JSON output
Almost all the closed-source LLMs now have an option to give JSON as output but what about open-source models? Some cloud providers like Groq have added the option to return JSON.
But what other options do we have? Let's add a strict type checking for JSON output.
We need 3 libraries for this
- Pydantic: For creating the JSON structure that we are expecting.
- Langchain_core: For parsing the output in the format defined in Pydantic
- Tenacity: For Retrying if parsing fails
Let’s try to understand this with an example:
Pydantic
$ pip install pydantic
from pydantic import BaseModel
class LLMOutput(BaseModel):
summary: str # Just a string
suggestions: list[str] # A list containing strings
The summary is of type string and suggestions contain a list of strings
JsonParser
$ pip install langchain_core
import json
from langchain_core.output_parsers import JsonOutputParser
class JsonParser:
def __init__(self, pydantic_object) -> None:
self._parser = JsonOutputParser(pydantic_object=pydantic_object)
def _basic_parser(self, content: str) -> json:
try:
return json.loads(content)
except Exception as e:
raise e
def parse(self, content: str) -> json:
parsed_output = self._parser.parse(content)
return parsed_output
Let’s see the combined code
import json
from pydantic import BaseModel
from langchain_core.output_parsers import JsonOutputParser
class LLMOutput(BaseModel):
summary: str # Just a string
suggestions: list[str] # A list containing strings
class JsonParser:
def __init__(self, pydantic_object) -> None:
self._parser = JsonOutputParser(pydantic_object=pydantic_object)
def _basic_parser(self, content: str) -> json:
try:
return json.loads(content)
except Exception as e:
raise e
def parse(self, content: str) -> json:
parsed_output = self._parser.parse(content)
return parsed_output
json_parser = JsonParser(LLMOutput)
json_parser.parse(<model_output_here>)
Tenacity
Now, we need to create a function that will take some input, process it through LLM, parse the output and return the output
$ pip install tenacity
from tenacity import stop, stop_after_retry, retry
@retry(stop=stop_after_retry(2))
def generate_content(input_query: str) -> dict:
# You need to replace this part with your llm call
model_output = LLM(input_query)
parsed_output = json_parser.parse(model_output)
return parsed_output
What we have done is very simple but very effective.
If the model doesn’t give the expected output; the parser will raise an exception. Tenacity will rerun the function, if even that fails then another retry and if that fails too then it will give an error.
Parsers ensure that we only get the output in the expected format.
Let’s now see the complete code in action:
import json
from pydantic import BaseModel
from langchain_core.output_parsers import JsonOutputParser
from tenacity import stop, stop_after_retry, retry
class LLMOutput(BaseModel):
summary: str # Just a string
suggestions: list[str] # A list containing strings
@retry(stop=stop_after_retry(2))
def generate_content(input_query: str) -> dict:
# You need to replace this part with your llm call
model_output = LLM(input_query)
parsed_output = json_parser.parse(model_output)
return parsed_output
class JsonParser:
def __init__(self, pydantic_object) -> None:
self._parser = JsonOutputParser(pydantic_object=pydantic_object)
def _basic_parser(self, content: str) -> json:
try:
return json.loads(content)
except Exception as e:
raise e
def parse(self, content: str) -> json:
parsed_output = self._parser.parse(content)
return parsed_output
json_parser = JsonParser(LLMOutput)
model_output = generate_content(<prompt_here with instructions to give only json output in defined structure>)
parsed_output = json_parser.parse(<model_output_here>)
print(parsed_output)
One problem with the proposed approach is retries increase the cost and time to generate an output.
A couple of libraries offer similar functionality out of the box but if you want more control then it’s always better to write your own code.