File size: 3,340 Bytes
ab6600f 3c83920 ab6600f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | from typing import Any, Optional
from smolagents.tools import Tool
import os
import urllib
import json
class SpecStreamGetRows(Tool):
name = "specstream_get_rows"
description = "Reads the results of a SpecStream run by find_id, each value with the URL of the page it came from. Also says whether the run is still going. Reading costs nothing. Needs SPECSTREAM_API_KEY."
inputs = {'find_id': {'type': 'string', 'description': 'The find_id returned by specstream_find_web_data.'}, 'limit': {'type': 'integer', 'description': 'How many rows to read, 1 to 100. Default 25.', 'nullable': True}, 'cursor': {'type': 'string', 'description': 'The next_cursor from a previous call, to read the next page.', 'nullable': True}}
output_type = "string"
def forward(self, find_id: str, limit: int | None = None, cursor: str | None = None) -> str:
import json
import os
from urllib.error import HTTPError, URLError
from urllib.parse import quote as urlquote, urlencode
from urllib.request import Request, urlopen
key = os.environ.get("SPECSTREAM_API_KEY", "").strip()
if not key:
return "Set SPECSTREAM_API_KEY to a SpecStream API key (https://specstream.ai/workspace?utm_source=smolagents)."
base = os.environ.get("SPECSTREAM_API_BASE", "https://api.specstream.ai").rstrip("/")
size = 25 if limit is None else max(1, min(int(limit), 100))
query = {"limit": str(size)}
if cursor:
query["cursor"] = cursor
path = "/v1/finds/%s/rows?%s" % (urlquote(find_id, safe=""), urlencode(query))
req = Request(
base + path,
headers={
"Authorization": "Bearer " + key,
"Accept": "application/json",
"User-Agent": "specstream-smolagents/0.1",
},
)
try:
with urlopen(req, timeout=120) as resp:
payload = json.loads(resp.read().decode("utf-8"))
except HTTPError as exc:
return "SpecStream answered HTTP %d for that find_id." % exc.code
except (URLError, OSError):
return "SpecStream could not be reached."
page = payload.get("data") if isinstance(payload, dict) else None
if not isinstance(page, dict):
return "SpecStream returned no rows page."
rows = page.get("rows") or []
lines = ["status: %s, accepted rows: %s" % (page.get("status"), page.get("accepted_count"))]
for number, row in enumerate(rows, 1):
parts = []
for cell in row.get("cells") or []:
shown = cell.get("display")
if shown in (None, ""):
continue
source = (cell.get("provenance") or {}).get("source_url")
parts.append("%s: %s%s" % (cell.get("label"), shown, " (source: %s)" % source if source else ""))
lines.append("%d. %s" % (number, "; ".join(parts)))
if not rows:
lines.append("No rows yet. If the status is still running, call again in a minute.")
cursor_next = page.get("next_cursor")
if cursor_next:
lines.append("next_cursor: %s" % cursor_next)
return "\n".join(lines)
def __init__(self, *args, **kwargs):
self.is_initialized = False
|