app.bettersight.io/backend/core/extractor.py

130 lines
4.6 KiB
Python

import os
import re
import json
import logging
from openai import OpenAI
import sentry_sdk
logger = logging.getLogger(__name__)
OPENROUTER_API_KEY = os.getenv('OPENROUTER_API_KEY')
FREE_MODELS = [m.strip() for m in os.getenv('OPENROUTER_FREE_MODELS', '').split(',') if m.strip()]
FALLBACK_MODEL = os.getenv('OPENROUTER_FALLBACK_MODEL', 'anthropic/claude-haiku-4-5')
class ExtractorError(Exception):
"""Raised when every OpenRouter model (free tier + paid fallback) has failed."""
pass
class RateLimitError(Exception):
"""Raised internally when an OpenRouter model returns a 429 / rate-limit response."""
pass
def _parse_json_response(text, model, label):
"""
Parses a model's raw text response into JSON via progressively more
aggressive repairs — ported unchanged from the legacy app.py
extraction logic: raw parse → regex object extraction →
trailing-comma/unquoted-key repair.
Data flow:
raw model text → strip code fences →
json.loads() (fast path) →
regex-extract the outermost {...} block and retry →
strip trailing commas + quote bare keys and retry →
parsed dict returned, or ValueError if all three fail
"""
cleaned = re.sub(r'```json|```', '', text).strip()
try:
return json.loads(cleaned)
except json.JSONDecodeError:
pass
match = re.search(r'\{[\s\S]*\}', cleaned)
if match:
try:
return json.loads(match.group(0))
except json.JSONDecodeError:
pass
try:
fixed = re.sub(r',\s*}', '}', cleaned)
fixed = re.sub(r',\s*]', ']', fixed)
fixed = re.sub(r'([{,])\s*([a-zA-Z0-9_]+)\s*:', r'\1"\2":', fixed)
return json.loads(fixed)
except json.JSONDecodeError:
raise ValueError(f'Could not parse response from {model} for {label}')
def _call_openrouter(model, prompt, content):
"""
Issues one chat completion request to OpenRouter for `model`.
`content` is appended to `prompt` when present — most callers
(via industries.adventure_travel.prompt.build_prompt) already embed
page content directly into `prompt`, so `content` is typically
empty; analysis_service's refresh path passes it explicitly too,
which is harmless (the model just sees the page text once via the
prompt and, redundantly, once more here).
Data flow:
model + prompt + content → OpenAI client (OpenRouter base_url) →
chat.completions.create() → raise RateLimitError on 429/rate-limit
responses (caller moves to the next model) → raw text →
_parse_json_response() → parsed dict returned
"""
client = OpenAI(api_key=OPENROUTER_API_KEY, base_url='https://openrouter.ai/api/v1')
message = f'{prompt}\n\n{content}' if content else prompt
try:
response = client.chat.completions.create(
model=model,
messages=[{'role': 'user', 'content': message}],
temperature=0.1,
max_tokens=4096
)
except Exception as e:
error_str = str(e)
if '429' in error_str or 'rate_limit' in error_str.lower():
raise RateLimitError(error_str)
raise
text = response.choices[0].message.content
return _parse_json_response(text, model, 'extraction')
def extract_with_openrouter(prompt, content=''):
"""
Cycles through free OpenRouter models in sequence. Falls back to
the paid fallback model if all free models fail or rate-limit.
Never returns None — raises ExtractorError only after every model
(free + paid) has failed.
Data flow:
prompt + content → for each model in FREE_MODELS + [FALLBACK_MODEL]:
_call_openrouter() → success: return parsed JSON
(logging + a Sentry breadcrumb if the paid fallback had to be used) →
RateLimitError: log, try next model →
any other error: log, try next model →
all models exhausted → raise ExtractorError
"""
models_to_try = FREE_MODELS + [FALLBACK_MODEL]
for model in models_to_try:
try:
result = _call_openrouter(model, prompt, content)
if model == FALLBACK_MODEL and FREE_MODELS:
logger.warning(f'Free models exhausted — used paid fallback: {model}')
sentry_sdk.capture_message('OpenRouter paid fallback used', level='warning')
return result
except RateLimitError:
logger.warning(f'Rate limit on {model} — trying next')
continue
except Exception as e:
logger.error(f'Model {model} failed: {str(e)}')
continue
raise ExtractorError('All OpenRouter models exhausted including paid fallback')