diffbot 3.0.0__py3-none-any.whl

diffbot/cli/__init__.py

@@ -0,0 +1,430 @@
+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 (
+    __version__,
+    CrawlEventType,
+    AuthError,
+    ExtractionError,
+    APIError,
+    ValidationError,
+    json_schema_format,
+)
+
+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()
+@click.version_option(__version__, "-V", "--version", prog_name="db")
+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")
+@click.option(
+    "--schema",
+    "schema_path",
+    type=click.Path(exists=True, dir_okay=False),
+    help="JSON Schema file constraining the output (implies --json)",
+)
+def ask(prompt: str, output: str = None, as_json: bool = False, schema_path: str = None):
+    """Ask a question to the Diffbot LLM"""
+    db = get_client()
+    stdin_content = sys.stdin.read() if not sys.stdin.isatty() else None
+
+    response_format = None
+    if schema_path:
+        as_json = True
+        try:
+            with open(schema_path) as f:
+                response_format = json_schema_format(json.load(f))
+        except json.JSONDecodeError as e:
+            click.echo(f"Error: {schema_path} is not valid JSON: {e}", err=True)
+            raise click.Abort()
+
+    # Without --schema, response_format stays None and ask_json picks its own
+    # permissive default; --json alone must not become {"type": "json_object"}.
+    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 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:
+                if as_json:
+                    response = json.dumps(
+                        db.ask_json(messages, response_format=response_format), indent=2
+                    )
+                else:
+                    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:
+                sys.stdout.write(
+                    json.dumps(db.ask_json(messages, response_format=response_format), indent=2)
+                )
+            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 ValidationError as e:
+        click.echo(f"Error: {e}", 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,21 @@
+import click
+
+from diffbot import Diffbot, resolve_token
+from diffbot._auth import CREDENTIALS_PATH, TOKEN_ENV_VAR
+
+
+def get_client() -> Diffbot:
+    """Build a Diffbot client using the shared credential resolution chain.
+
+    Looks at the DIFFBOT_API_TOKEN env var, then ~/.diffbot/credentials.
+    """
+    token = resolve_token()
+    if not token:
+        click.echo(
+            "Error: no Diffbot API token found.\n"
+            f"  Set a {TOKEN_ENV_VAR} environment variable, or\n"
+            f"  write '{TOKEN_ENV_VAR}=YOUR_TOKEN' to {CREDENTIALS_PATH}",
+            err=True,
+        )
+        raise click.Abort()
+    return Diffbot(token=token)