Simple Hack to always generate JSON from LLM

Trick to always get JSON output

Photo by Bernd 📷 Dittrich on Unsplash

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

1. Pydantic: For creating the JSON structure that we are expecting.
2. Langchain_core: For parsing the output in the format defined in Pydantic
3. 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.