diffbot-python 0.1.0__py3-none-any.whl

diffbot/__init__.py

@@ -0,0 +1,29 @@
+"""
+diffbot - Python client library for the Diffbot APIs.
+"""
+
+__version__ = "0.1.0"
+
+from .client import Diffbot, DiffbotAsync
+from .crawl import CrawlEvent, CrawlEventType
+from .errors import (
+    APIError,
+    AuthError,
+    DiffbotError,
+    ExtractionError,
+    RateLimitError,
+    ValidationError,
+)
+
+__all__ = [
+    "Diffbot",
+    "DiffbotAsync",
+    "CrawlEvent",
+    "CrawlEventType",
+    "DiffbotError",
+    "AuthError",
+    "ExtractionError",
+    "RateLimitError",
+    "APIError",
+    "ValidationError",
+]

diffbot/ask.py

@@ -0,0 +1,48 @@
+"""Diffbot LLM RAG API: stream a chat completion."""
+
+import json
+from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, List
+
+if TYPE_CHECKING:
+    from .client import Diffbot, DiffbotAsync
+
+
+def _build_payload(client: Any, messages: List[Dict[str, str]]) -> tuple:
+    headers = {"Authorization": f"Bearer {client.token}"}
+    payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
+    return headers, payload
+
+
+def _parse_chunk(line: str):
+    try:
+        chunk = json.loads(line.replace("data: ", ""))
+    except json.JSONDecodeError:
+        return None
+    choices = chunk.get("choices")
+    if choices and choices[0].get("delta", {}).get("content"):
+        return choices[0]["delta"]["content"]
+    return None
+
+
+def ask(client: "Diffbot", messages: List[Dict[str, str]]) -> Iterator[str]:
+    headers = {"Authorization": f"Bearer {client.token}"}
+    payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
+    with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
+        client._raise_for_status(response)
+        for line in response.iter_lines():
+            if line:
+                content = _parse_chunk(line)
+                if content:
+                    yield content
+
+
+async def ask_async(client: "DiffbotAsync", messages: List[Dict[str, str]]) -> AsyncIterator[str]:
+    headers = {"Authorization": f"Bearer {client.token}"}
+    payload = {"model": "diffbot-small-xl", "messages": messages, "stream": True}
+    async with client._http.stream("POST", client.llm_url, headers=headers, json=payload) as response:
+        client._raise_for_status(response)
+        async for line in response.aiter_lines():
+            if line:
+                content = _parse_chunk(line)
+                if content:
+                    yield content

diffbot/cli/__init__.py

@@ -0,0 +1,399 @@
+import json
+import re
+import sys
+import textwrap
+from email.utils import parsedate_to_datetime
+
+import click
+from rich.console import Console
+from rich.live import Live
+from rich.markdown import Markdown
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from datetime import datetime
+
+from diffbot import CrawlEventType, AuthError, ExtractionError, APIError
+
+from ._common import get_client
+
+console = Console()
+is_interactive = sys.stdout.isatty()
+
+
+def print_markdown(text):
+    if is_interactive:
+        console.print(Markdown(text))
+    else:
+        sys.stdout.write(text)
+        sys.stdout.write("\n")
+
+
+@click.group()
+def main():
+    """
+    Diffbot 🤖 Structure the world's knowledge.
+    """
+    pass
+
+
+@main.command(name="extract")
+@click.argument("url")
+@click.option("-o", "--output", type=click.Path(), help="Write output to file instead of stdout")
+@click.option("-f", "--format", "fmt", type=click.Choice(["markdown", "json"]), default="markdown", help="Output format")
+@click.option("-a", "--api", type=str, default="analyze", help="Diffbot API to use")
+def extract(url, output, fmt, api):
+    """Get structured content from a URL."""
+    db = get_client()
+    try:
+        if is_interactive and not output:
+            with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+                progress.add_task(description="Extracting content...", total=None)
+                data = db.extract(url, api=api, fmt=fmt)
+        else:
+            data = db.extract(url, api=api, fmt=fmt)
+
+        if fmt == "json":
+            output_str = json.dumps(data, indent=2)
+        else:
+            objects = data.get("objects", [])
+            if not objects:
+                output_str = json.dumps(data, indent=2)
+            else:
+                obj = objects[0]
+                title = obj.get("title", "")
+                url = obj.get("pageUrl", "")
+                content = obj.get("content", "")
+                output_str = f"Title: {title}\n\nURL: {url}\n\nContent: {content}"
+
+        if output:
+            with open(output, "w") as f:
+                f.write(output_str)
+            click.echo(f"Output written to {output}")
+        else:
+            print_markdown(output_str)
+    except ExtractionError as e:
+        click.echo(f"Extraction error {e.error_code}: {e.error}", err=True)
+        raise click.Abort()
+    except AuthError:
+        click.echo("Error: Invalid or unauthorized API token.", err=True)
+        raise click.Abort()
+    except APIError as e:
+        click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
+        raise click.Abort()
+
+
+@main.command()
+@click.argument("prompt")
+@click.option("-o", "--output", type=click.Path(), help="Write output to file instead of stdout")
+@click.option("--json", "as_json", is_flag=True, help="Output from LLM as JSON")
+def ask(prompt: str, output: str = None, as_json: bool = False):
+    """Ask a question to the Diffbot LLM"""
+    db = get_client()
+    stdin_content = sys.stdin.read() if not sys.stdin.isatty() else None
+    interactive_mode = is_interactive and not output and not as_json
+
+    messages = [
+        {
+            "role": "system",
+            "content": f"Current local time: {datetime.now().strftime('%A, %B %d, %Y, %I:%M:%S %p %Z')}",
+        },
+        {"role": "user", "content": prompt},
+    ]
+
+    if as_json:
+        messages[1]["content"] += "\nReturn the output as a JSON object. Do not include any other text outside of the JSON object."
+
+    if stdin_content:
+        messages[1]["content"] = f"<input>{stdin_content}</input>\n" + messages[1]["content"]
+
+    try:
+        if interactive_mode:
+            response_text = ""
+            with Live("", auto_refresh=False, console=console) as live:
+                for chunk in db.ask(messages):
+                    response_text += chunk
+                    live.update(Markdown(response_text))
+                    live.refresh()
+
+            messages.append({"role": "assistant", "content": response_text})
+
+            while True:
+                try:
+                    sys.stdout.write("\n> ")
+                    sys.stdout.flush()
+                    user_input = input()
+                    messages.append({"role": "user", "content": user_input})
+                except (KeyboardInterrupt, EOFError):
+                    console.print("\nExiting chat...")
+                    return
+
+                response_text = ""
+                with Live("", auto_refresh=False, console=console) as live:
+                    for chunk in db.ask(messages):
+                        response_text += chunk
+                        live.update(Markdown(response_text))
+                        live.refresh()
+                messages.append({"role": "assistant", "content": response_text})
+        else:
+            response = ""
+            if output:
+                for chunk in db.ask(messages):
+                    response += chunk
+                with open(output, "w") as f:
+                    f.write(response)
+                click.echo(f"Output written to {output}")
+            elif as_json:
+                buffer = ""
+                for chunk in db.ask(messages):
+                    buffer += chunk
+                buffer = buffer[buffer.find("{") : buffer.rfind("}") + 1]
+                sys.stdout.write(buffer)
+            else:
+                for chunk in db.ask(messages):
+                    sys.stdout.write(chunk)
+    except (KeyboardInterrupt, EOFError):
+        console.print("\nExiting chat...")
+    except AuthError:
+        click.echo("Error: Invalid or unauthorized API token.", err=True)
+        raise click.Abort()
+    except APIError as e:
+        click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
+        raise click.Abort()
+
+
+@main.command()
+@click.argument("site")
+@click.option("--hops", type=int, default=2, help="Maximum link depth from seed URLs")
+@click.option("--job-name", type=str, help="Name for the crawler job (generated if not provided)")
+@click.option("--max-to-crawl", type=int, default=100, help="Maximum number of pages to crawl")
+@click.option("--max-to-process", type=int, default=100, help="Maximum number of pages to process")
+@click.option("--restrict-domain", is_flag=True, default=True, help="Restrict crawling to the same domain as seeds")
+@click.option("--api-url", type=str, default="", help="Diffbot API endpoint to use for processing")
+@click.option("--crawl-delay", type=float, default=-1, help="Delay between requests to the same domain in seconds")
+@click.option("--url-crawl-pattern", type=str, help="Only crawl URLs matching this pattern")
+@click.option("--url-process-pattern", type=str, help="Only process URLs matching this pattern")
+@click.option("--obey-robots", is_flag=True, help="Obey robots.txt rules")
+@click.option("--use-proxies", is_flag=True, help="Use proxies for crawling")
+@click.option("--custom-headers", type=str, help="Custom HTTP headers to send with requests (newline separated)")
+@click.option("-o", "--output", type=click.Path(), help="Write output to file instead of stdout")
+@click.option("-f", "--format", "fmt", type=click.Choice(["markdown", "json"]), default="markdown", help="Output format")
+def crawl(site, hops, job_name, max_to_crawl, max_to_process, restrict_domain,
+          api_url, crawl_delay, url_crawl_pattern, url_process_pattern, obey_robots,
+          use_proxies, custom_headers, output, fmt):
+    """Crawl a website using the Diffbot Crawler API."""
+    db = get_client()
+    events = []
+    try:
+        with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+            progress.add_task(description=f"Crawling {site}...", total=None)
+
+            for event in db.crawl(
+                site,
+                hops=hops,
+                job_name=job_name,
+                max_to_crawl=max_to_crawl,
+                max_to_process=max_to_process,
+                restrict_domain=restrict_domain,
+                api_url=api_url,
+                crawl_delay=crawl_delay,
+                url_crawl_pattern=url_crawl_pattern,
+                url_process_pattern=url_process_pattern,
+                obey_robots=obey_robots,
+                use_proxies=use_proxies,
+                custom_headers=custom_headers,
+                watch=True,
+            ):
+                events.append(event)
+                if event.event_type == CrawlEventType.JOB_CREATED:
+                    progress.console.print(f"Job created: {event.details['job_name']}")
+                elif event.event_type == CrawlEventType.URL_PROCESSED:
+                    status = event.details.get("status", "")
+                    icon = "✓" if status == "Success" else "!"
+                    progress.console.print(f"  [{icon}] {event.details['url']}")
+
+        if fmt == "json":
+            output_str = json.dumps(
+                [{"event_type": e.event_type.value, "timestamp": e.timestamp, "details": e.details} for e in events],
+                indent=2,
+            )
+        else:
+            lines = []
+            for event in events:
+                if event.event_type == CrawlEventType.JOB_CREATED:
+                    lines.append(f"# Job: {event.details['job_name']}\n")
+                elif event.event_type == CrawlEventType.URL_PROCESSED:
+                    lines.append(f"- [{event.details.get('status', 'unknown')}] {event.details['url']}")
+            output_str = "\n".join(lines)
+
+        if output:
+            with open(output, "w") as f:
+                f.write(output_str)
+            click.echo(f"Output written to {output}")
+        else:
+            print_markdown(output_str)
+    except AuthError:
+        click.echo("Error: Invalid or unauthorized API token.", err=True)
+        raise click.Abort()
+    except APIError as e:
+        click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
+        raise click.Abort()
+
+
+@main.command()
+@click.argument("job_name", required=False)
+def crawl_list_jobs(job_name):
+    """List Diffbot crawler jobs or get details for a specific job."""
+    db = get_client()
+    try:
+        if job_name:
+            job = db.crawl_get_job(job_name)
+            print_markdown(f"# Job: {job_name}\n\n```json\n{json.dumps(job, indent=2)}\n```")
+        else:
+            jobs = db.crawl_list_jobs()
+            if not jobs:
+                print_markdown("No crawler jobs found.")
+                return
+            lines = ["# Diffbot Crawler Jobs\n"]
+            for job in jobs:
+                name = job.get("name", "Unknown")
+                job_type = job.get("type", "Unknown")
+                status = job.get("jobStatus", {}).get("message", "Unknown")
+                lines.append(f"## {name}")
+                lines.append(f"Type: {job_type}")
+                lines.append(f"Status: {status}")
+                if "pageCrawlSuccesses" in job:
+                    lines.append(f"Pages Crawled: {job['pageCrawlSuccesses']}")
+                if "pageProcessSuccesses" in job:
+                    lines.append(f"Pages Processed: {job['pageProcessSuccesses']}")
+                lines.append("\n---\n")
+            print_markdown("\n".join(lines))
+    except AuthError:
+        click.echo("Error: Invalid or unauthorized API token.", err=True)
+        raise click.Abort()
+    except APIError as e:
+        click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
+        raise click.Abort()
+
+
+@main.command()
+@click.argument("job_name")
+def crawl_delete_job(job_name):
+    """Delete a Diffbot crawler job."""
+    db = get_client()
+    try:
+        with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+            progress.add_task(description=f"Deleting job {job_name}...", total=None)
+            db.crawl_delete_job(job_name)
+            progress.console.print(f"Job {job_name} deleted.")
+    except AuthError:
+        click.echo("Error: Invalid or unauthorized API token.", err=True)
+        raise click.Abort()
+    except APIError as e:
+        click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
+        raise click.Abort()
+
+
+def _format_date(date_str: str) -> str:
+    if not date_str:
+        return ""
+    try:
+        dt = parsedate_to_datetime(date_str)
+        return f"{dt.strftime('%b')} {dt.day}, {dt.year}"
+    except Exception:
+        return date_str[:12].strip()
+
+
+def _score_color(score: float) -> str:
+    if score > 0.85:
+        return "bright_green"
+    elif score >= 0.7:
+        return "green"
+    elif score >= 0.5:
+        return "yellow"
+    return "red"
+
+
+def _strip_markdown(text: str) -> str:
+    text = re.sub(r'!\[.*?\]\(.*?\)', '', text)          # images
+    text = re.sub(r'\[([^\]]*)\]\([^)]*\)', r'\1', text) # links → label
+    text = re.sub(r'^#{1,6}\s+', '', text, flags=re.MULTILINE)  # headings
+    text = re.sub(r'^\s*[-*+]\s+', '', text, flags=re.MULTILINE) # bullets
+    text = re.sub(r'^\s*\d+\.\s+', '', text, flags=re.MULTILINE) # numbered lists
+    text = re.sub(r'[`*_~]{1,3}', '', text)              # inline code/bold/italic
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
+
+
+@main.command(name="web-search")
+@click.argument("query")
+@click.option("-n", "--num-results", type=int, default=None, help="Number of results to return")
+@click.option("-m", "--max-tokens", type=int, default=None, help="Limit total response tokens (for agentic use cases)")
+@click.option("-f", "--format", "fmt", type=click.Choice(["list", "json", "text"]), default="list", help="Output format (text is plain/agent-friendly)")
+def web_search(query: str, num_results, max_tokens, fmt: str):
+    """Search the web using the Diffbot LLM web search API."""
+    db = get_client()
+    try:
+        if is_interactive:
+            with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+                progress.add_task(description="Searching...", total=None)
+                data = db.web_search(query, num_results=num_results, max_tokens=max_tokens)
+        else:
+            data = db.web_search(query, num_results=num_results, max_tokens=max_tokens)
+
+        if fmt == "json":
+            sys.stdout.write(json.dumps(data, indent=2))
+            sys.stdout.write("\n")
+            return
+
+        results = data.get("search_results", [])
+
+        if fmt == "text":
+            for i, r in enumerate(results, 1):
+                date_fmt = _format_date(r.get("date", ""))
+                content = " ".join(r.get("content", "").split())
+                lines = [
+                    f"[{i}] {r.get('title', '')} (score: {r.get('score', 0):.3f})",
+                    f"URL: {r.get('pageUrl', '')}",
+                ]
+                if date_fmt:
+                    lines.append(f"Date: {date_fmt}")
+                lines += [f"Content: {content}", "---"]
+                sys.stdout.write("\n".join(lines) + "\n")
+            return
+
+        time_ms = data.get("timeMs", 0)
+
+        for i, r in enumerate(results, 1):
+            score = r.get("score", 0)
+            title = r.get("title", "")
+            url = r.get("pageUrl", "")
+            date = r.get("date", "")
+            content = _strip_markdown(r.get("content", ""))
+
+            date_fmt = _format_date(date)
+            color = _score_color(score)
+            prefix_len = len(date_fmt) + 3 if date_fmt else 0  # " — "
+            snippet_width = max(80, console.width * 2 - prefix_len)
+            snippet = textwrap.shorten(content, width=snippet_width, placeholder="...")
+
+            console.print(f"[bold]#{i}: {title}[/bold] [[{color}]{score:.3f}[/{color}]]")
+            console.print(f"[dim blue]{url}[/dim blue]")
+            if date_fmt:
+                console.print(f"[dim]{date_fmt}[/dim] — {snippet}")
+            else:
+                console.print(snippet)
+            console.print()
+
+        console.print(f"[dim]{len(results)} result(s) in {time_ms}ms[/dim]")
+    except AuthError:
+        click.echo("Error: Invalid or unauthorized API token.", err=True)
+        raise click.Abort()
+    except APIError as e:
+        click.echo(f"API error {e.status_code}: {e.message or e.body}", err=True)
+        raise click.Abort()
+
+
+from .dql import dql as _dql_group  # noqa: E402
+from .entities import entities as _entities_cmd  # noqa: E402
+
+main.add_command(_dql_group)
+main.add_command(_entities_cmd)

diffbot/cli/__main__.py

@@ -0,0 +1,4 @@
+from . import main
+
+if __name__ == "__main__":
+    main()

diffbot/cli/_common.py

@@ -0,0 +1,36 @@
+import os
+import pathlib
+
+import click
+
+from diffbot import Diffbot
+
+CREDENTIALS_PATH = pathlib.Path.home() / ".diffbot" / "credentials"
+
+
+def resolve_token() -> str:
+    """Return the Diffbot API token from the env var, falling back to ~/.diffbot/credentials."""
+    token = os.environ.get("DIFFBOT_API_TOKEN", "").strip()
+    if token:
+        return token
+
+    if CREDENTIALS_PATH.exists():
+        for line in CREDENTIALS_PATH.read_text().splitlines():
+            line = line.strip()
+            if line.startswith("DIFFBOT_API_TOKEN="):
+                return line[len("DIFFBOT_API_TOKEN="):].strip()
+
+    return ""
+
+
+def get_client() -> Diffbot:
+    token = resolve_token()
+    if not token:
+        click.echo(
+            "Error: no Diffbot API token found.\n"
+            "  Set a DIFFBOT_API_TOKEN environment variable, or\n"
+            f"  write 'DIFFBOT_API_TOKEN=YOUR_TOKEN' to {CREDENTIALS_PATH}",
+            err=True,
+        )
+        raise click.Abort()
+    return Diffbot(token=token)