diffbot-python 0.1.0__py3-none-any.whl

diffbot/client.py

@@ -0,0 +1,285 @@
+"""The Diffbot client classes (sync and async)."""
+
+import pathlib
+from types import TracebackType
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Sequence, Type, Union
+
+import httpx
+
+from . import __version__
+from .errors import APIError, AuthError, RateLimitError, ValidationError
+from .extract import extract as _extract, extract_async as _extract_async
+from .ask import ask as _ask, ask_async as _ask_async
+from .crawl import (
+    CrawlEvent,
+    crawl as _crawl,
+    crawl_async as _crawl_async,
+    crawl_delete_job as _crawl_delete_job,
+    crawl_delete_job_async as _crawl_delete_job_async,
+    crawl_get_job as _crawl_get_job,
+    crawl_get_job_async as _crawl_get_job_async,
+    crawl_list_jobs as _crawl_list_jobs,
+    crawl_list_jobs_async as _crawl_list_jobs_async,
+)
+from .kg import (
+    dql as _dql,
+    dql_async as _dql_async,
+    dql_parallel as _dql_parallel,
+    dql_refresh_ontology as _dql_refresh_ontology,
+)
+from .web_search import (
+    WEB_SEARCH_BASE,
+    web_search as _web_search,
+    web_search_async as _web_search_async,
+)
+from .nlp import (
+    NLP_BASE,
+    entities as _entities,
+    entities_async as _entities_async,
+)
+
+EXTRACT_BASE = "https://api.diffbot.com/v3"
+CRAWL_BASE = "https://api.diffbot.com/v3/crawl"
+DIFFBOT_LLM_BASE = "https://llm.diffbot.com/rag/v1/chat/completions"
+DEFAULT_TIMEOUT = 30.0
+
+
+class Diffbot:
+    """Client for the Diffbot APIs.
+
+    Example:
+        >>> from diffbot import Diffbot
+        >>> db = Diffbot(token=os.getenv("DIFFBOT_API_TOKEN"))
+        >>> db.extract("https://example.com")
+    """
+
+    def __init__(
+        self,
+        token: str,
+        *,
+        timeout: float = DEFAULT_TIMEOUT,
+        analyze_url: str = EXTRACT_BASE,
+        llm_url: str = DIFFBOT_LLM_BASE,
+        crawler_url: str = CRAWL_BASE,
+        web_search_url: str = WEB_SEARCH_BASE,
+        nlp_url: str = NLP_BASE,
+        transport: Optional[httpx.BaseTransport] = None,
+    ):
+        if not token:
+            raise ValidationError("token is required")
+        self.token = token
+        self.analyze_url = analyze_url
+        self.llm_url = llm_url
+        self.crawler_url = crawler_url
+        self.web_search_url = web_search_url
+        self.nlp_url = nlp_url
+        self._http = httpx.Client(
+            timeout=timeout,
+            headers={"User-Agent": f"diffbot-python/{__version__}"},
+            transport=transport,
+        )
+
+    def __enter__(self) -> "Diffbot":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._http.close()
+
+    def _raise_for_status(self, response: httpx.Response) -> None:
+        if response.is_success:
+            return
+        status = response.status_code
+        body = response.text
+        if status in (401, 403):
+            raise AuthError(status, body)
+        if status == 429:
+            raise RateLimitError(status, body, retry_after=response.headers.get("retry-after"))
+        raise APIError(status, body)
+
+    def extract(self, url: str, api: str = "analyze", fmt: str = "markdown") -> Dict[str, Any]:
+        """Extract structured content from a URL. Returns the raw Diffbot API response."""
+        return _extract(self, url, api=api, fmt=fmt)
+
+    def ask(self, messages: List[Dict[str, str]]) -> Iterator[str]:
+        """Stream a response from the Diffbot LLM RAG API."""
+        yield from _ask(self, messages)
+
+    def crawl(self, site: str, **kwargs: Any) -> Iterator[CrawlEvent]:
+        """Start a crawl job."""
+        yield from _crawl(self, site, **kwargs)
+
+    def crawl_list_jobs(self) -> List[Dict[str, Any]]:
+        """List all crawler jobs for this token."""
+        return _crawl_list_jobs(self)
+
+    def crawl_get_job(self, job_name: str) -> Dict[str, Any]:
+        """Get the status of a crawler job."""
+        return _crawl_get_job(self, job_name)
+
+    def crawl_delete_job(self, job_name: str) -> None:
+        """Delete a crawler job."""
+        _crawl_delete_job(self, job_name)
+
+    def dql(
+        self,
+        query: str,
+        *,
+        size: int = 10,
+        from_: int = 0,
+        format: str = "json",
+        filter: Optional[str] = None,
+        exportspec: Optional[str] = None,
+        extra: Optional[Dict[str, str]] = None,
+        raw: bool = False,
+    ) -> Union[Dict[str, Any], bytes]:
+        """Run a DQL query against the Diffbot Knowledge Graph.
+
+        Returns the parsed JSON response, or the raw response bytes when raw=True
+        (e.g. to retrieve CSV/export formats undecoded).
+        """
+        return _dql(self, query, size=size, from_=from_, format=format, filter=filter, exportspec=exportspec, extra=extra, raw=raw)
+
+    def dql_parallel(self, queries: Sequence[Dict[str, Any]], *, workers: int = 8) -> List[Union[Dict[str, Any], bytes]]:
+        """Run multiple DQL queries concurrently. Each item is a dict of dql() keyword args."""
+        return _dql_parallel(self, queries, workers=workers)
+
+    def dql_refresh_ontology(self, dest: pathlib.Path) -> None:
+        """Download the Diffbot Knowledge Graph ontology and write it to dest."""
+        _dql_refresh_ontology(self, dest)
+
+    def web_search(self, text: str, *, num_results: Optional[int] = None, max_tokens: Optional[int] = None) -> Dict[str, Any]:
+        """Search the web via the Diffbot LLM web search API."""
+        return _web_search(self, text, num_results=num_results, max_tokens=max_tokens)
+
+    def entities(self, text: str, *, lang: str = "auto") -> Dict[str, Any]:
+        """Identify and resolve entities and sentiment in text using the Diffbot NLP API.
+
+        Entity IDs can be looked up in the Knowledge Graph via dql() using
+        id:or("id1","id2","id3") — no type: declaration required.
+        """
+        return _entities(self, text, lang=lang)
+
+
+class DiffbotAsync:
+    """Async client for Diffbot APIs.
+
+    Example:
+        >>> async with DiffbotAsync(token=os.getenv("DIFFBOT_API_TOKEN")) as db:
+        ...     result = await db.extract("https://example.com")
+    """
+
+    def __init__(
+        self,
+        token: str,
+        *,
+        timeout: float = DEFAULT_TIMEOUT,
+        analyze_url: str = EXTRACT_BASE,
+        llm_url: str = DIFFBOT_LLM_BASE,
+        crawler_url: str = CRAWL_BASE,
+        web_search_url: str = WEB_SEARCH_BASE,
+        nlp_url: str = NLP_BASE,
+        transport: Optional[httpx.AsyncBaseTransport] = None,
+    ):
+        if not token:
+            raise ValidationError("token is required")
+        self.token = token
+        self.analyze_url = analyze_url
+        self.llm_url = llm_url
+        self.crawler_url = crawler_url
+        self.web_search_url = web_search_url
+        self.nlp_url = nlp_url
+        self._http = httpx.AsyncClient(
+            timeout=timeout,
+            headers={"User-Agent": f"diffbot-python/{__version__}"},
+            transport=transport,
+        )
+
+    async def __aenter__(self) -> "DiffbotAsync":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
+        await self.close()
+
+    async def close(self) -> None:
+        await self._http.aclose()
+
+    def _raise_for_status(self, response: httpx.Response) -> None:
+        if response.is_success:
+            return
+        status = response.status_code
+        body = response.text
+        if status in (401, 403):
+            raise AuthError(status, body)
+        if status == 429:
+            raise RateLimitError(status, body, retry_after=response.headers.get("retry-after"))
+        raise APIError(status, body)
+
+    async def extract(self, url: str, api: str = "analyze", fmt: str = "markdown") -> Dict[str, Any]:
+        """Extract structured content from a URL. Returns the raw Diffbot API response."""
+        return await _extract_async(self, url, api=api, fmt=fmt)
+
+    async def ask(self, messages: List[Dict[str, str]]) -> AsyncIterator[str]:
+        """Stream a response from the Diffbot LLM RAG API."""
+        async for chunk in _ask_async(self, messages):
+            yield chunk
+
+    async def crawl(self, site: str, **kwargs: Any) -> AsyncIterator[CrawlEvent]:
+        """Start a crawl job. Pass watch=True to poll until completion and yield URL_PROCESSED events."""
+        async for event in _crawl_async(self, site, **kwargs):
+            yield event
+
+    async def crawl_list_jobs(self) -> List[Dict[str, Any]]:
+        """List all crawler jobs for this token."""
+        return await _crawl_list_jobs_async(self)
+
+    async def crawl_get_job(self, job_name: str) -> Dict[str, Any]:
+        """Get the status of a crawler job."""
+        return await _crawl_get_job_async(self, job_name)
+
+    async def crawl_delete_job(self, job_name: str) -> None:
+        """Delete a crawler job."""
+        await _crawl_delete_job_async(self, job_name)
+
+    async def dql(
+        self,
+        query: str,
+        *,
+        size: int = 10,
+        from_: int = 0,
+        format: str = "json",
+        filter: Optional[str] = None,
+        exportspec: Optional[str] = None,
+        extra: Optional[Dict[str, str]] = None,
+        raw: bool = False,
+    ) -> Union[Dict[str, Any], bytes]:
+        """Run a DQL query against the Diffbot Knowledge Graph.
+
+        Returns the parsed JSON response, or the raw response bytes when raw=True
+        (e.g. to retrieve CSV/export formats undecoded).
+        """
+        return await _dql_async(self, query, size=size, from_=from_, format=format, filter=filter, exportspec=exportspec, extra=extra, raw=raw)
+
+    async def web_search(self, text: str, *, num_results: Optional[int] = None, max_tokens: Optional[int] = None) -> Dict[str, Any]:
+        """Search the web via the Diffbot LLM web search API."""
+        return await _web_search_async(self, text, num_results=num_results, max_tokens=max_tokens)
+
+    async def entities(self, text: str, *, lang: str = "auto") -> Dict[str, Any]:
+        """Identify and resolve entities and sentiment in text using the Diffbot NLP API.
+
+        Entity IDs can be looked up in the Knowledge Graph via dql() using
+        id:or("id1","id2","id3") — no type: declaration required.
+        """
+        return await _entities_async(self, text, lang=lang)

diffbot/crawl.py

@@ -0,0 +1,270 @@
+"""Diffbot Crawler API: start crawls, manage crawler jobs."""
+
+import asyncio
+import csv
+import io
+import time
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List, Optional
+
+from .errors import APIError
+
+if TYPE_CHECKING:
+    from .client import Diffbot, DiffbotAsync
+
+
+class CrawlEventType(Enum):
+    JOB_CREATED = "job_created"
+    URL_PROCESSED = "url_processed"
+
+
+@dataclass
+class CrawlEvent:
+    """An event yielded while a crawl is in flight."""
+    event_type: CrawlEventType
+    timestamp: str
+    details: Dict[str, Any]
+
+
+def crawl(
+    client: "Diffbot",
+    site: str,
+    hops: int = 2,
+    job_name: Optional[str] = None,
+    max_to_crawl: int = 100,
+    max_to_process: int = 100,
+    restrict_domain: bool = True,
+    api_url: str = "",
+    crawl_delay: float = -1,
+    url_crawl_pattern: Optional[str] = None,
+    url_process_pattern: Optional[str] = None,
+    obey_robots: bool = False,
+    use_proxies: bool = False,
+    custom_headers: Optional[str] = None,
+    watch: bool = False,
+    poll_interval: float = 2.0,
+) -> Iterator[CrawlEvent]:
+    """Create a crawler job and yield a JOB_CREATED event. If watch=True, poll until
+    completion and yield URL_PROCESSED events for each crawled URL."""
+    if not site.startswith("http"):
+        site = f"https://{site}"
+    if not job_name:
+        job_name = f"crawl-{int(time.time())}"
+
+    params = _build_crawl_params(
+        client, job_name, site, max_to_crawl, max_to_process, hops, restrict_domain,
+        api_url, crawl_delay, url_crawl_pattern, url_process_pattern,
+        obey_robots, use_proxies, custom_headers,
+    )
+    response = client._http.get(client.crawler_url, params=params)
+    client._raise_for_status(response)
+
+    yield CrawlEvent(
+        event_type=CrawlEventType.JOB_CREATED,
+        timestamp=datetime.now().isoformat(),
+        details={"job_name": job_name},
+    )
+
+    if not watch:
+        return
+
+    seen_urls: set = set()
+
+    while True:
+        status_response = client._http.get(
+            client.crawler_url,
+            params={"token": client.token, "name": job_name},
+        )
+        client._raise_for_status(status_response)
+        jobs = status_response.json().get("jobs", [])
+        if not jobs:
+            break
+        job_status = jobs[0].get("jobStatus", {})
+        status_code = job_status.get("status", 0)
+
+        urls_response = client._http.get(
+            f"{client.crawler_url}/data",
+            params={"token": client.token, "name": job_name, "type": "urls"},
+            follow_redirects=True,
+        )
+        client._raise_for_status(urls_response)
+        if urls_response.content:
+            for event in _parse_url_csv(urls_response.text, seen_urls):
+                yield event
+
+        if status_code not in (0, 7):
+            message = job_status.get("message", "")
+            if "fail" in message.lower() or "error" in message.lower():
+                raise APIError(500, f"Crawler job failed: {message}")
+            break
+
+        time.sleep(poll_interval)
+
+
+def _build_crawl_params(client: Any, job_name: str, site: str, max_to_crawl: int,
+                        max_to_process: int, hops: int, restrict_domain: bool,
+                        api_url: str, crawl_delay: float, url_crawl_pattern: Optional[str],
+                        url_process_pattern: Optional[str], obey_robots: bool,
+                        use_proxies: bool, custom_headers: Optional[str]) -> Dict[str, Any]:
+    params: Dict[str, Any] = {
+        "token": client.token,
+        "name": job_name,
+        "seeds": site,
+        "maxToCrawl": max_to_crawl,
+        "maxToProcess": max_to_process,
+        "maxHops": hops,
+        "restrictDomain": 1 if restrict_domain else 0,
+    }
+    if api_url:
+        params["apiUrl"] = api_url
+    if crawl_delay > 0:
+        params["crawlDelay"] = crawl_delay
+    if url_crawl_pattern:
+        params["urlCrawlPattern"] = url_crawl_pattern
+    if url_process_pattern:
+        params["urlProcessPattern"] = url_process_pattern
+    if obey_robots:
+        params["obeyRobots"] = 1
+    if use_proxies:
+        params["useProxies"] = 1
+    if custom_headers:
+        params["customHeaders"] = custom_headers
+    return params
+
+
+def _parse_url_csv(text: str, seen_urls: set) -> List[CrawlEvent]:
+    events = []
+    reader = csv.DictReader(io.StringIO(text))
+    for row in reader:
+        url = row.get("Url", "").strip('"')
+        if url and url not in seen_urls:
+            seen_urls.add(url)
+            events.append(CrawlEvent(
+                event_type=CrawlEventType.URL_PROCESSED,
+                timestamp=row.get("Crawled Time", datetime.now().isoformat()),
+                details={"url": url, "status": row.get("Crawl Status", "unknown")},
+            ))
+    return events
+
+
+async def crawl_async(
+    client: "DiffbotAsync",
+    site: str,
+    hops: int = 2,
+    job_name: Optional[str] = None,
+    max_to_crawl: int = 100,
+    max_to_process: int = 100,
+    restrict_domain: bool = True,
+    api_url: str = "",
+    crawl_delay: float = -1,
+    url_crawl_pattern: Optional[str] = None,
+    url_process_pattern: Optional[str] = None,
+    obey_robots: bool = False,
+    use_proxies: bool = False,
+    custom_headers: Optional[str] = None,
+    watch: bool = False,
+    poll_interval: float = 2.0,
+) -> AsyncIterator[CrawlEvent]:
+    if not site.startswith("http"):
+        site = f"https://{site}"
+    if not job_name:
+        job_name = f"crawl-{int(time.time())}"
+
+    params = _build_crawl_params(
+        client, job_name, site, max_to_crawl, max_to_process, hops, restrict_domain,
+        api_url, crawl_delay, url_crawl_pattern, url_process_pattern,
+        obey_robots, use_proxies, custom_headers,
+    )
+    response = await client._http.get(client.crawler_url, params=params)
+    client._raise_for_status(response)
+
+    yield CrawlEvent(
+        event_type=CrawlEventType.JOB_CREATED,
+        timestamp=datetime.now().isoformat(),
+        details={"job_name": job_name},
+    )
+
+    if not watch:
+        return
+
+    seen_urls: set = set()
+
+    while True:
+        status_response = await client._http.get(
+            client.crawler_url,
+            params={"token": client.token, "name": job_name},
+        )
+        client._raise_for_status(status_response)
+        jobs = status_response.json().get("jobs", [])
+        if not jobs:
+            break
+        job_status = jobs[0].get("jobStatus", {})
+        status_code = job_status.get("status", 0)
+
+        urls_response = await client._http.get(
+            f"{client.crawler_url}/data",
+            params={"token": client.token, "name": job_name, "type": "urls"},
+            follow_redirects=True,
+        )
+        client._raise_for_status(urls_response)
+        if urls_response.content:
+            for event in _parse_url_csv(urls_response.text, seen_urls):
+                yield event
+
+        if status_code not in (0, 7):
+            message = job_status.get("message", "")
+            if "fail" in message.lower() or "error" in message.lower():
+                raise APIError(500, f"Crawler job failed: {message}")
+            break
+
+        await asyncio.sleep(poll_interval)
+
+
+async def crawl_list_jobs_async(client: "DiffbotAsync") -> List[Dict[str, Any]]:
+    response = await client._http.get(client.crawler_url, params={"token": client.token})
+    client._raise_for_status(response)
+    return response.json().get("jobs", [])
+
+
+async def crawl_get_job_async(client: "DiffbotAsync", job_name: str) -> Dict[str, Any]:
+    response = await client._http.get(
+        client.crawler_url,
+        params={"token": client.token, "name": job_name},
+    )
+    client._raise_for_status(response)
+    jobs = response.json().get("jobs", [])
+    return jobs[0] if jobs else {}
+
+
+async def crawl_delete_job_async(client: "DiffbotAsync", job_name: str) -> None:
+    response = await client._http.get(
+        client.crawler_url,
+        params={"token": client.token, "name": job_name, "delete": 1},
+    )
+    client._raise_for_status(response)
+
+
+def crawl_list_jobs(client: "Diffbot") -> List[Dict[str, Any]]:
+    response = client._http.get(client.crawler_url, params={"token": client.token})
+    client._raise_for_status(response)
+    return response.json().get("jobs", [])
+
+
+def crawl_get_job(client: "Diffbot", job_name: str) -> Dict[str, Any]:
+    response = client._http.get(
+        client.crawler_url,
+        params={"token": client.token, "name": job_name},
+    )
+    client._raise_for_status(response)
+    jobs = response.json().get("jobs", [])
+    return jobs[0] if jobs else {}
+
+
+def crawl_delete_job(client: "Diffbot", job_name: str) -> None:
+    response = client._http.get(
+        client.crawler_url,
+        params={"token": client.token, "name": job_name, "delete": 1},
+    )
+    client._raise_for_status(response)

diffbot/errors.py

@@ -0,0 +1,51 @@
+"""Exception hierarchy for the Diffbot SDK."""
+
+import json
+from typing import Optional
+
+
+class DiffbotError(Exception):
+    """Base class for all Diffbot SDK errors."""
+
+
+class ValidationError(DiffbotError):
+    """Client-side validation failed (e.g. missing token, malformed argument)."""
+
+
+class APIError(DiffbotError):
+    """The Diffbot API returned an error response."""
+
+    def __init__(self, status_code: int, body: str):
+        self.status_code = status_code
+        self.body = body
+        self.message: Optional[str] = None
+        self.request_id: Optional[str] = None
+        try:
+            data = json.loads(body)
+            self.message = data.get("message")
+            self.request_id = data.get("requestId")
+        except (ValueError, AttributeError):
+            pass
+        display = self.message or (body[:200] + ("..." if len(body) > 200 else ""))
+        super().__init__(f"Diffbot API error {status_code}: {display}")
+
+
+class AuthError(APIError):
+    """Authentication failed (401, 403)."""
+
+
+class RateLimitError(APIError):
+    """Rate limit exceeded (429)."""
+
+    def __init__(self, status_code: int, body: str, retry_after: Optional[str] = None):
+        super().__init__(status_code, body)
+        self.retry_after = retry_after
+
+
+class ExtractionError(DiffbotError):
+    """The Diffbot API returned a 200 but reported an extraction failure."""
+
+    def __init__(self, error_code: int, error: str):
+        self.error_code = error_code
+        self.error = error
+        super().__init__(f"Diffbot extraction error {error_code}: {error}")

diffbot/extract.py

@@ -0,0 +1,45 @@
+"""Diffbot Analyze API: extract structured content from a URL."""
+
+from typing import TYPE_CHECKING, Any, Dict
+
+if TYPE_CHECKING:
+    from .client import Diffbot, DiffbotAsync
+
+from .errors import ExtractionError
+
+
+def _build_params(client: Any, url: str, fmt: str) -> Dict[str, Any]:
+    params = {"token": client.token, "url": url, "timeout": 30000}
+    if fmt == "markdown":
+        params["mode"] = "llm"
+    return params
+
+
+def _parse_response(client: Any, data: Dict[str, Any]) -> Dict[str, Any]:
+    if "errorCode" in data:
+        raise ExtractionError(data["errorCode"], data.get("error", ""))
+    return data
+
+
+def _normalize_url(url: str) -> str:
+    return url if url.startswith("http") else f"https://{url}"
+
+
+def extract(client: "Diffbot", url: str, api: str = "analyze", fmt: str = "markdown") -> Dict[str, Any]:
+    url = _normalize_url(url)
+    response = client._http.get(
+        f"{client.analyze_url}/{api}",
+        params=_build_params(client, url, fmt),
+    )
+    client._raise_for_status(response)
+    return _parse_response(client, response.json())
+
+
+async def extract_async(client: "DiffbotAsync", url: str, api: str = "analyze", fmt: str = "markdown") -> Dict[str, Any]:
+    url = _normalize_url(url)
+    response = await client._http.get(
+        f"{client.analyze_url}/{api}",
+        params=_build_params(client, url, fmt),
+    )
+    client._raise_for_status(response)
+    return _parse_response(client, response.json())