Using JSON mode with Text Gen endpoints

Ensure Text Gen outputs fit into your desired JSON schema.

OctoAIs Large Language Models (LLMs) can generate generate outputs that not only adhere to JSON format but also align with your unique schema specifications.

Supported models (Updated September 5, 2024 5PM PT)

  • Llama 3.1 8B
  • Llama 3.1 70B
  • Hermes 2 Pro Llama 3 8B
  • Mistral 7B
  • Nous Hermes Mixtral 8x7B
  • Mixtral 8x7B
  • WizardLM 8x22B

Getting started

Setup credentials:

$export OCTOAI_TOKEN=YOUR_TOKEN_HERE

Curl example (Mistral-7B): Let’s say that you want to ensure that your LLM responses format user feedback about cars into a usable JSON format. To do so, you provide the LLM with a reponse schema ensuring that it knows it must provide “color” and “maker” in a structured format—see “response format below”:

$curl -X POST "https://text.octoai.run/v1/chat/completions" \
>  -H "Content-Type: application/json" \
>  -H "Authorization: Bearer $OCTOAI_TOKEN" \
>  --data-raw '{
>        "messages": [
>            {
>                "role": "system",
>                "content": "You are a helpful assistant."
>            },
>            {
>                "role": "user",
>                "content": "the car was black and it was a toyota camry."
>            }
>        ],
>        "model": "mistral-7b-instruct",
>        "max_tokens": 512,
>        "presence_penalty": 0,
>        "temperature": 0.1,
>        "top_p": 0.9,
>        "response_format": {
>            "type": "json_object",
>            "schema": {"properties": {"color": {"title": "Color", "type": "string"}, "maker": {"title": "Maker", "type": "string"}}, "required": ["color", "maker"], "title": "Car", "type": "object"}
>        }
>    }'

The LLM will respond in the exact schema specified:

${
>  "id": "chatcmpl-d5d81b7c80b249ea8177f95f68a51d8e",
>  "object": "chat.completion",
>  "created": 1709830931,
>  "model": "mistral-7b-instruct",
>  "choices": [
>    {
>      "index": 0,
>      "message": {
>        "role": "assistant",
>        "content": "{\"color\": \"black\", \"maker\": \"Toyota”, \"}",
>        "function_call": null
>      },
>      "finish_reason": "stop"
>    }
>  ],
>  "usage": {
>    "prompt_tokens": 98,
>    "completion_tokens": 16,
>    "total_tokens": 114
>  }
>}

Pydantic and OctoAI’s Python SDK

Pydantic is a popular Python library for data validation and settings management using Python type annotations. By combining Pydantic with the OctoAI SDK, you can easily define the desired JSON schema for your LLM responses and ensure that the generated content adheres to that structure.

First, make sure you have the required packages installed:

$python3 -m pip install openai pydantic==2.5.3

Basic example

Let’s start with a basic example to demonstrate how Pydantic and the OctoAI SDK work together. In this example, we’ll define a simple Car model with color and maker attributes, and ask the LLM to generate a response that fits this schema.

1from octoai.client import OctoAI
2from octoai.text_gen import ChatCompletionResponseFormat, ChatMessage
3from pydantic import BaseModel, Field
4from typing import List
5
6client = OctoAI()
7
8class Car(BaseModel):
9    color: str
10    maker: str
11
12completion = client.text_gen.create_chat_completion(
13    model="mistral-7b-instruct",
14    messages=[
15        ChatMessage(role="system", content="You are a helpful assistant."),
16        ChatMessage(role="user", content="the car was black and it was a toyota camry."),
17    ],
18    max_tokens=512,
19    presence_penalty=0,
20    temperature=0.1,
21    top_p=0.9,
22    response_format=ChatCompletionResponseFormat(
23        type="json_object",
24        schema=Car.model_json_schema(),
25    ),
26)
27
28print(completion.choices[0].message.content)

The key points to note here are:

  1. We import the necessary classes from the OctoAI SDK: Client, TextModel, and ChatCompletionResponseFormat.

  2. We define a Car class inheriting from BaseModel, specifying the color and maker attributes with their expected types.

  3. When creating the chat completion, we set the response_format using ChatCompletionResponseFormat and include the JSON schema generated from our Car model using Car.model_json_schema().

The output will be a JSON object adhering to the specified schema:

1{ "color": "black", "maker": "Toyota" }

Array example

Next, let’s look at an example involving arrays. Suppose we want the LLM to generate a list of names based on a given prompt. We can define a Meeting model with a names attribute of type List[str].

1from octoai.client import OctoAI
2from octoai.text_gen import ChatCompletionResponseFormat, ChatMessage
3from pydantic import BaseModel, Field
4from typing import List
5
6client = OctoAI()
7
8class Meeting(BaseModel):
9    names: List[str]
10
11
12chat_completion = client.text_gen.create_chat_completion(
13    model="<model>",
14    messages=[
15        ChatMessage(role="system", content="You are a helpful assistant."),
16        ChatMessage(role="user", content="John and Jane meet the day after"),
17    ],
18    temperature=0,
19    response_format=ChatCompletionResponseFormat(
20        type="json_object",
21        schema=Meeting.model_json_schema()
22    ),
23)
24
25print(chat_completion.choices[0].message.content)

The LLM will generate a response containing an array of names:

1{ "names": ["John", "Jane"] }

Nested example

Finally, let’s explore a more complex example involving nested models. In this case, we’ll define a Person model with name and age attributes, and a Result model containing a sorted list of Person objects.

1class Person(BaseModel):
2    """The object representing a person with name and age"""
3
4    name: str = Field(description="Name of the person")
5    age: int = Field(description="The age of the person")
6
7
8class Result(BaseModel):
9    """The format of the answer."""
10
11    sorted_list: List[Person] = Field(description="List of the sorted objects")
12
13
14completion = octoai.text_gen.create_chat_completion(
15    model="mistral-7b-instruct",
16    messages=[
17        {
18            "role": "system",
19            "content": "You are a helpful assistant designed to output JSON.",
20        },
21        {
22            "role": "user",
23            "content": "Alice is 10 years old, Bob is 7 and carol is 2. Sort them by age in ascending order.",
24        },
25    ],
26    max_tokens=512,
27    presence_penalty=0,
28    temperature=0.1,
29    top_p=0.9,
30    response_format=ChatCompletionResponseFormat(
31        type="json_object",
32        schema=Result.model_json_schema(),
33    ),
34)
35
36print(completion.choices[0].message.content)

In this example:

  1. We define a Person model with name and age attributes, along with descriptions using the Field function from Pydantic.
  2. We define a Result model containing a sorted_list attribute of type List[Person].
  3. When creating the chat completion, we set the response_format using ChatCompletionResponseFormat and include the JSON schema generated from our Result model.

The LLM will generate a response containing a sorted list of Person objects:

1{
2  "sorted_list": [
3    { "name": "Carol", "age": 2 },
4    { "name": "Bob", "age": 7 },
5    { "name": "Alice", "age": 10 }
6  ]
7}

Instructor

Instructor makes it easy to reliably get structured data like JSON from Large Language Models (LLMs). Read more here

Install

$python3 -m pip install instructor

Example

1import os
2import openai
3from pydantic import BaseModel
4import instructor
5
6client = openai.OpenAI(
7    base_url="https://text.octoai.run/v1",
8    api_key=os.environ["OCTOAI_TOKEN"],
9)
10
11
12# By default, the patch function will patch the ChatCompletion.create and ChatCompletion.create methods to support the response_model parameter
13client = instructor.patch(client, mode=instructor.Mode.JSON_SCHEMA)
14
15
16# Now, we can use the response_model parameter using only a base model
17# rather than having to use the OpenAISchema class
18class UserExtract(BaseModel):
19    name: str
20    age: int
21
22
23user: UserExtract = client.chat.completions.create(
24    model="mistral-7b-instruct",
25    response_model=UserExtract,
26    messages=[
27        {"role": "user", "content": "Extract jason is 25 years old"},
28    ],
29)
30
31print(user.model_dump_json(indent=2))

Let’s break down the code step by step:

After importing the necessary modules and setting the clients, we:

  1. We use the instructor.patch function to patch the ChatCompletion.create method of the OctoAI client. This allows us to use the response_model parameter directly with a Pydantic model.

  2. We define a Pydantic model called UserExtract that represents the desired structure of the extracted user information. In this case, it has two fields: name (a string) and age (an integer).

  3. We call the chat.completions.create method of the patched OctoAI client, specifying the model (mistral-7b-instruct), the response_model (our UserExtract model), and the user message that contains the information we want to extract.

  4. Finally, we print the extracted user information using the model_dump_json method, which serializes the Pydantic model to a JSON string with indentation for better readability.

The output will be a JSON object containing the extracted user information, adhering to the specified UserExtract schema:

1{
2  "name": "jason",
3  "age": 25
4}

By leveraging Instructor and the OctoAI SDK, you can easily define the desired output schema and ensure that the LLM generates structured data that fits your application’s requirements. This simplifies the process of integrating LLM-generated content into your software systems.