Pythonjson_extraction.py
Copy
1
import os
2
import json
3
import csv
4
from typing import Optional
5
import openai
6
from pydantic import BaseModel, Field
7
 
8
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
9
 
10
# ── 1. JSON Mode (valid JSON, no schema guarantee) ──────────────────
11
def extract_product_json_mode(description: str) -> dict:
12
    response = client.chat.completions.create(
13
        model="gpt-4o-mini",
14
        messages=[
15
            {
16
                "role": "system",
17
                "content": """Extract product info as JSON with keys:
18
                name (string), price (number), category (string), 
19
                features (array of strings), in_stock (boolean)""",
20
            },
21
            {"role": "user", "content": description},
22
        ],
23
        response_format={"type": "json_object"},
24
        temperature=0,
25
    )
26
    return json.loads(response.choices[0].message.content)
27
 
28
# ── 2. Structured Outputs (schema guaranteed) ───────────────────────
29
class ProductFeature(BaseModel):
30
    name: str
31
    value: str
32
 
33
class Product(BaseModel):
34
    name: str
35
    brand: Optional[str] = None
36
    price_usd: Optional[float] = None
37
    category: str
38
    features: list[str]
39
    in_stock: bool
40
    rating: Optional[float] = Field(None, ge=0.0, le=5.0)
41
 
42
def extract_product_structured(description: str) -> Product:
43
    response = client.beta.chat.completions.parse(
44
        model="gpt-4o-mini",
45
        messages=[
46
            {
47
                "role": "system",
48
                "content": "Extract product information from the description.",
49
            },
50
            {"role": "user", "content": description},
51
        ],
52
        response_format=Product,
53
        temperature=0,
54
    )
55
    return response.choices[0].message.parsed
56
 
57
# ── 3. Batch extraction → CSV ───────────────────────────────────────
58
PRODUCT_DESCRIPTIONS = [
59
    "Sony WH-1000XM5 wireless headphones by Sony. $349.99. Active noise cancelling, 30hr battery, foldable design. In stock. Rated 4.8/5.",
60
    "The Kindle Paperwhite (Amazon) offers a 6.8-inch display, waterproof design, 3 months battery life. $139.99. Currently out of stock.",
61
    "Apple AirPods Pro (2nd gen) with H2 chip, adaptive transparency, MagSafe charging case. $249. In stock. 4.7 stars.",
62
]
63
 
64
print("=== Batch Extraction ===")
65
products = []
66
for desc in PRODUCT_DESCRIPTIONS:
67
    try:
68
        product = extract_product_structured(desc)
69
        products.append(product)
70
        print(f"  ✓ {product.name} — ${product.price_usd} ({product.category})")
71
    except Exception as e:
72
        print(f"  ✗ Extraction failed: {e}")
73
 
74
# Write to CSV
75
with open("products.csv", "w", newline="") as f:
76
    writer = csv.DictWriter(f, fieldnames=["name", "brand", "price_usd", "category", "in_stock", "rating"])
77
    writer.writeheader()
78
    for p in products:
79
        writer.writerow({
80
            "name": p.name, "brand": p.brand, "price_usd": p.price_usd,
81
            "category": p.category, "in_stock": p.in_stock, "rating": p.rating,
82
        })
83
print(f"Written {len(products)} products to products.csv")
84