Ross
92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
import os
|
|
import json
|
|
import time
|
|
from typing import Optional, Dict, Any
|
|
|
|
try:
|
|
import openai
|
|
except Exception:
|
|
openai = None
|
|
|
|
|
|
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
|
|
DEFAULT_MODEL = os.environ.get('CARDS_LLM_MODEL', 'gpt-4o-mini')
|
|
|
|
|
|
PROMPT_TEMPLATE = (
|
|
"""
|
|
You are a JSON generator. Given a dinosaur name, return EXACTLY one JSON object matching the schema below.
|
|
Only return JSON (no explanatory text).
|
|
|
|
Schema:
|
|
{
|
|
"name": string,
|
|
"speed_kmh": number|null, // numeric speed in km/h
|
|
"weight_kg": number|null, // numeric weight in kg
|
|
"height_m": number|null, // numeric height in meters
|
|
"intelligence_score": integer|null, // 1-5 scale or null
|
|
"facts": [string,...],
|
|
"sources": [{"title":string,"url":string},...]
|
|
}
|
|
|
|
Rules:
|
|
- Use units: km/h, kg, m.
|
|
- If you cannot find a reliable number, use null.
|
|
- Provide up to 5 short facts.
|
|
- Include at least one source object with a URL when available.
|
|
Name: {name}
|
|
"""
|
|
)
|
|
|
|
|
|
def _ensure_openai():
|
|
if openai is None:
|
|
raise RuntimeError('openai package not installed; add openai to requirements')
|
|
if not OPENAI_API_KEY:
|
|
raise RuntimeError('OPENAI_API_KEY not set in environment')
|
|
openai.api_key = OPENAI_API_KEY
|
|
|
|
|
|
def _extract_json(text: str) -> Optional[Dict[str, Any]]:
|
|
text = text.strip()
|
|
# try direct parse
|
|
try:
|
|
return json.loads(text)
|
|
except Exception:
|
|
pass
|
|
# heuristic: find first { ... } block
|
|
start = text.find('{')
|
|
end = text.rfind('}')
|
|
if start != -1 and end != -1 and end > start:
|
|
try:
|
|
return json.loads(text[start:end+1])
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def fetch_dinosaur_structured(name: str, model: Optional[str] = None, max_retries: int = 2) -> Dict[str, Any]:
|
|
"""Call the configured LLM and return parsed JSON per schema.
|
|
|
|
Raises RuntimeError on misconfiguration or ValueError if parsing fails.
|
|
"""
|
|
model = model or DEFAULT_MODEL
|
|
if openai is None:
|
|
raise RuntimeError('openai package not available')
|
|
_ensure_openai()
|
|
|
|
prompt = PROMPT_TEMPLATE.format(name=name)
|
|
for attempt in range(max_retries + 1):
|
|
resp = openai.ChatCompletion.create(
|
|
model=model,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
temperature=0.0,
|
|
max_tokens=600,
|
|
)
|
|
content = resp['choices'][0]['message']['content'].strip()
|
|
data = _extract_json(content)
|
|
if data:
|
|
return data
|
|
time.sleep(1 + attempt)
|
|
raise ValueError('Failed to parse LLM response as JSON')
|