bhushanct
511e36b660
After each capture the controller asks the vision model to locate a 'Next' control (returned as normalized screen coordinates), clicks it via OS input, scrolls to top, and continues the loop. Configurable via SCREEN_LEADS_AUTO_NEXT / _NEXT_LOAD_PAUSE / _MAX_AUTO_NEXT, with a per-run safety cap. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
"""Vision extraction via Claude.
|
||
|
||
Two calls:
|
||
- extract_lead(): screenshots + a recipe's schema -> raw structured JSON.
|
||
- classify_page(): one screenshot -> {site, page_type, address bar url}, used
|
||
only when no native browser URL is available (visual site detection).
|
||
|
||
Both use structured outputs (output_config.format) so the response is
|
||
schema-validated JSON, not free text.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import json
|
||
from functools import lru_cache
|
||
from typing import Any
|
||
|
||
from config import ANTHROPIC_MODEL, MAX_IMAGES_PER_CALL
|
||
from ai.recipes.base import SiteRecipe
|
||
|
||
|
||
@lru_cache(maxsize=1)
|
||
def _client():
|
||
import anthropic
|
||
|
||
# Resolves ANTHROPIC_API_KEY (or an `ant auth login` profile) from the env.
|
||
return anthropic.Anthropic()
|
||
|
||
|
||
def _image_block(path: str) -> dict[str, Any]:
|
||
with open(path, "rb") as f:
|
||
data = base64.standard_b64encode(f.read()).decode("utf-8")
|
||
return {
|
||
"type": "image",
|
||
"source": {"type": "base64", "media_type": "image/png", "data": data},
|
||
}
|
||
|
||
|
||
def _first_json(response) -> dict[str, Any]:
|
||
for block in response.content:
|
||
if block.type == "text":
|
||
return json.loads(block.text)
|
||
raise RuntimeError("model returned no text block")
|
||
|
||
|
||
def extract_lead(recipe: SiteRecipe, screenshot_paths: list[str]) -> dict[str, Any]:
|
||
paths = screenshot_paths[:MAX_IMAGES_PER_CALL]
|
||
content: list[dict[str, Any]] = [_image_block(p) for p in paths]
|
||
content.append({"type": "text", "text": recipe.extraction_prompt})
|
||
|
||
response = _client().messages.create(
|
||
model=ANTHROPIC_MODEL,
|
||
max_tokens=2048,
|
||
messages=[{"role": "user", "content": content}],
|
||
output_config={"format": {"type": "json_schema", "schema": recipe.extraction_schema}},
|
||
)
|
||
return _first_json(response)
|
||
|
||
|
||
_CLASSIFY_SCHEMA = {
|
||
"type": "object",
|
||
"properties": {
|
||
"site": {"type": "string"}, # e.g. "linkedin", "other"
|
||
"page_type": {"type": "string"}, # e.g. "profile", "search", "feed", "other"
|
||
"address_bar_url": {"type": "string"},
|
||
},
|
||
"required": ["site", "page_type", "address_bar_url"],
|
||
"additionalProperties": False,
|
||
}
|
||
|
||
_CLASSIFY_PROMPT = (
|
||
"This is a screenshot of a web browser. Identify the site and page from the "
|
||
"browser's address bar and the visible layout. Return: 'site' (lowercase, "
|
||
"e.g. 'linkedin', or 'other'); 'page_type' ('profile' for a single person's "
|
||
"page, otherwise 'search', 'feed', 'messaging', or 'other'); and "
|
||
"'address_bar_url' (the URL text you can read, or \"\")."
|
||
)
|
||
|
||
|
||
def classify_page(screenshot_path: str) -> dict[str, Any]:
|
||
response = _client().messages.create(
|
||
model=ANTHROPIC_MODEL,
|
||
max_tokens=512,
|
||
messages=[{
|
||
"role": "user",
|
||
"content": [_image_block(screenshot_path), {"type": "text", "text": _CLASSIFY_PROMPT}],
|
||
}],
|
||
output_config={"format": {"type": "json_schema", "schema": _CLASSIFY_SCHEMA}},
|
||
)
|
||
return _first_json(response)
|
||
|
||
|
||
_NEXT_SCHEMA = {
|
||
"type": "object",
|
||
"properties": {
|
||
"found": {"type": "boolean"},
|
||
"x": {"type": "number"}, # 0..1 fraction of image width (button centre)
|
||
"y": {"type": "number"}, # 0..1 fraction of image height
|
||
"label": {"type": "string"},
|
||
},
|
||
"required": ["found", "x", "y", "label"],
|
||
"additionalProperties": False,
|
||
}
|
||
|
||
_NEXT_PROMPT = (
|
||
"Look at this screenshot for a control that advances to the NEXT item or page — "
|
||
"for example a button or link labelled 'Next', 'Next result', 'See next profile', "
|
||
"or a right-facing pagination arrow ('>' / '›' / '→'). Ignore 'Back'/'Previous' "
|
||
"and unrelated arrows. If such a control is clearly visible, set found=true and give "
|
||
"its CENTRE position as fractions of the image: x = left→right (0.0–1.0), "
|
||
"y = top→bottom (0.0–1.0), plus its visible label. If none is visible, "
|
||
"return found=false, x=0, y=0, label=\"\"."
|
||
)
|
||
|
||
|
||
def find_next_button(screenshot_path: str) -> dict[str, Any]:
|
||
response = _client().messages.create(
|
||
model=ANTHROPIC_MODEL,
|
||
max_tokens=256,
|
||
messages=[{
|
||
"role": "user",
|
||
"content": [_image_block(screenshot_path), {"type": "text", "text": _NEXT_PROMPT}],
|
||
}],
|
||
output_config={"format": {"type": "json_schema", "schema": _NEXT_SCHEMA}},
|
||
)
|
||
return _first_json(response)
|