lm_extract_schema_from_json.py

import json
import os

import openai
from dotenv import load_dotenv

load_dotenv()

openai.api_key = os.environ.get("OPENAI_API_KEY")

def call_model(prompt):
  response = openai.Completion.create(
    model="text-davinci-003",
    prompt=prompt,
    temperature=0,
    max_tokens=2000,
  )
  return response.choices[0].text

def main():
  prompt = """
Extract a JSON schema that adheres to the structure of all of the following objects and is generic as possible.
Use `$ref`s where it makes sense.

Example object1:
{
  "id": 1,
  "name": "Smartphone",
  "description": "A high-end smartphone with cutting-edge features.",
  "price": 999.99,
  "brand": "Samsung",
  "category": "electronics",
  "tags": ["smartphone", "high-end", "Samsung"],
  "image": "https://example.com/smartphone.jpg",
  "variants": [
    {
      "id": 1,
      "name": "128GB",
      "price": 999.99,
      "image": "https://example.com/smartphone-128gb.jpg",
      "options": [
        {
          "name": "Color",
          "values": ["black", "silver"]
        },
        {
          "name": "Carrier",
          "values": ["Verizon", "AT&T", "T-Mobile"]
        }
      ]
    }
  ]
}

Example object 2:
{
  "id": 2,
  "name": "T-Shirt",
  "description": "A classic cotton t-shirt for everyday wear.",
  "price": 19.99,
  "brand": "Hanes",
  "category": "clothing",
  "tags": ["t-shirt", "cotton", "Hanes"],
  "image": "https://example.com/t-shirt.jpg",
  "variants": [
    {
      "id": 1,
      "name": "Small",
      "price": 19.99,
      "image": "https://example.com/t-shirt-small.jpg",
      "options": [
        {
          "name": "Color",
          "values": ["white", "black", "gray"]
        }
      ]
    },
    {
      "id": 2,
      "name": "Medium",
      "price": 19.99,
      "image": "https://example.com/t-shirt-medium.jpg",
      "options": [
        {
          "name": "Color",
          "values": ["white", "black", "gray"]
        }
      ]
    },
    {
      "id": 3,
      "name": "Large",
      "price": 19.99,
      "image": "https://example.com/t-shirt-large.jpg",
      "options": [
        {
          "name": "Color",
          "values": ["white", "black", "gray"]
        }
      ]
    }
  ]
}

Example object 3:
{
  "id": 3,
  "name": "The Great Gatsby",
  "description": "A classic novel by F. Scott Fitzgerald.",
  "price": 12.99,
  "brand": "Penguin Books",
  "category": "books",
  "tags": ["fiction", "classic", "F. Scott Fitzgerald"],
  "image": "https://example.com/great-gatsby.jpg",
  "variants": []
}


JSON Schema:
  """

  print(json.dumps(json.loads(call_model(prompt)), indent=2))

if __name__ == "__main__":
  main()