mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-20 11:42:23 +02:00
62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
"""
|
|
Schema-driven page extraction (the Stagehand `extract` port, on our terms).
|
|
|
|
One cheap aux call turns the current page's text into the JSON the agent asked
|
|
for, so the main model never spends context on 15k chars of raw page text. The
|
|
aux model is whatever cheap tier the user's provider offers (registry-resolved
|
|
by the caller); nothing here hardcodes a model.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
P_MAX_PAGE_CHARS = 12000
|
|
P_MAX_OUT_TOKENS = 1200
|
|
P_MAX_SCHEMA_CHARS = 2000
|
|
|
|
|
|
def first_json(text: str) -> str:
|
|
"""The model's output minus any prose/fences around the JSON, or ''."""
|
|
cleaned = re.sub(r"```(?:json)?|```", "", text or "")
|
|
m = re.search(r"\{.*\}|\[.*\]", cleaned, re.DOTALL)
|
|
if not m:
|
|
return ""
|
|
try:
|
|
return json.dumps(json.loads(m.group(0)), ensure_ascii=False)
|
|
except Exception:
|
|
return ""
|
|
|
|
|
|
async def extract_structured(
|
|
aux_client, aux_model, page_text: str, instruction: str, schema: dict | None = None,
|
|
) -> str:
|
|
"""Compact JSON string per the instruction (+ optional schema), or '' on
|
|
any miss so the caller can hand the agent an honest fallback."""
|
|
if not aux_client or not aux_model or not page_text:
|
|
return ""
|
|
shape = (
|
|
f"Return JSON matching this schema exactly:\n{json.dumps(schema)[:P_MAX_SCHEMA_CHARS]}"
|
|
if isinstance(schema, dict) and schema else "Return one compact JSON object."
|
|
)
|
|
prompt = (
|
|
"Extract data from this web page text.\n"
|
|
f"What to extract: {instruction}\n{shape}\n"
|
|
"Output ONLY the JSON, no prose, no code fences. Use only what is on the "
|
|
'page, never guess. If the requested data is not on the page, output '
|
|
'{"not_found": true, "reason": "<one short line>"}.\n\n'
|
|
f"PAGE TEXT:\n{page_text[:P_MAX_PAGE_CHARS]}"
|
|
)
|
|
try:
|
|
resp = await aux_client.messages.create(
|
|
model=aux_model, max_tokens=P_MAX_OUT_TOKENS,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
)
|
|
text = "".join(getattr(b, "text", "") for b in (resp.content or []))
|
|
return first_json(text)
|
|
except Exception as e:
|
|
logger.debug(f"[browser-extract] aux extraction failed: {e}")
|
|
return ""
|