Pythonconstrained_output.py
Copy
1
import os
2
from enum import Enum
3
from typing import Optional
4
import openai
5
from pydantic import BaseModel, Field
6
 
7
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
8
 
9
# ── Define Pydantic models ───────────────────────────────────────────
10
class Sentiment(str, Enum):
11
    POSITIVE = "positive"
12
    NEGATIVE = "negative"
13
    NEUTRAL = "neutral"
14
 
15
class EntityType(str, Enum):
16
    PERSON = "person"
17
    ORGANIZATION = "organization"
18
    LOCATION = "location"
19
    PRODUCT = "product"
20
 
21
class Entity(BaseModel):
22
    name: str
23
    type: EntityType
24
    confidence: float = Field(ge=0.0, le=1.0)
25
 
26
class DocumentAnalysis(BaseModel):
27
    title: str = Field(description="Short descriptive title for the document")
28
    summary: str = Field(max_length=200)
29
    sentiment: Sentiment
30
    key_entities: list[Entity]
31
    topics: list[str] = Field(max_length=5)
32
    requires_followup: bool
33
    priority: int = Field(ge=1, le=5, description="Priority 1-5, 5 is highest")
34
 
35
# ── Structured Outputs with Pydantic ────────────────────────────────
36
SAMPLE_TEXT = """
37
Apple Inc. announced today that CEO Tim Cook will present at the upcoming 
38
developer conference in San Francisco next month. The announcement follows 
39
disappointing Q3 earnings that missed analyst expectations by 8%. 
40
The stock dropped 4% on the news. Engineers in the Cupertino office are 
41
working on fixes to address customer complaints about the latest software update.
42
"""
43
 
44
def analyze_document(text: str) -> DocumentAnalysis:
45
    response = client.beta.chat.completions.parse(
46
        model="gpt-4o-mini",
47
        messages=[
48
            {
49
                "role": "system",
50
                "content": "You are a document analysis assistant. Extract structured information from text.",
51
            },
52
            {
53
                "role": "user",
54
                "content": f"Analyze this document:\n\n{text}",
55
            },
56
        ],
57
        response_format=DocumentAnalysis,
58
    )
59
    return response.choices[0].message.parsed
60
 
61
result = analyze_document(SAMPLE_TEXT)
62
print("=== Document Analysis ===")
63
print(f"Title: {result.title}")
64
print(f"Summary: {result.summary}")
65
print(f"Sentiment: {result.sentiment.value}")
66
print(f"Priority: {result.priority}/5")
67
print(f"Requires Follow-up: {result.requires_followup}")
68
print(f"Topics: {result.topics}")
69
print("\nEntities:")
70
for entity in result.key_entities:
71
    print(f"  {entity.name} ({entity.type.value}) — {entity.confidence:.0%} confidence")
72
 
73
# ── Batch extraction ────────────────────────────────────────────────
74
class ContactInfo(BaseModel):
75
    name: Optional[str] = None
76
    email: Optional[str] = None
77
    phone: Optional[str] = None
78
    company: Optional[str] = None
79
 
80
texts = [
81
    "Please contact John Smith at john.smith@acme.com or call 555-0123",
82
    "Reach out to Sarah Connor (Cyberdyne Systems) at s.connor@cyberdyne.io",
83
    "No contact info here — just some random text",
84
]
85
 
86
print("\n=== Batch Contact Extraction ===")
87
for text in texts:
88
    r = client.beta.chat.completions.parse(
89
        model="gpt-4o-mini",
90
        messages=[{"role": "user", "content": f"Extract contact info: {text}"}],
91
        response_format=ContactInfo,
92
    )
93
    contact = r.choices[0].message.parsed
94
    if contact.email or contact.phone:
95
        print(f"  Found: {contact.name} | {contact.email} | {contact.phone}")
96
    else:
97
        print(f"  No contact: {text[:40]}...")
98