diffbot-python 0.1.0__py3-none-any.whl

diffbot/cli/dql.py

@@ -0,0 +1,308 @@
+import csv
+import io
+import json
+import pathlib
+import re
+import shutil
+import time
+
+import click
+import sys
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.table import Table
+
+from diffbot import DiffbotError
+
+from . import ontology
+from ._common import get_client, resolve_token
+
+
+class _DqlGroup(click.Group):
+    _SECTIONS = [
+        (
+            "Commands",
+            ["export"],
+        ),
+        (
+            "Advanced Commands",
+            ["init", "ontology", "probe"],
+        ),
+    ]
+
+    def invoke(self, ctx):
+        try:
+            return super().invoke(ctx)
+        except (FileNotFoundError, KeyError, DiffbotError) as e:
+            raise click.ClickException(str(e))
+
+    def format_commands(self, ctx, formatter):
+        section_cmds = {label: [] for label, _ in self._SECTIONS}
+        uncategorized = []
+
+        for name in self.list_commands(ctx):
+            cmd = self.commands.get(name)
+            if cmd is None or cmd.hidden:
+                continue
+            placed = False
+            for label, names in self._SECTIONS:
+                if name in names:
+                    section_cmds[label].append((name, cmd))
+                    placed = True
+                    break
+            if not placed:
+                uncategorized.append((name, cmd))
+
+        for label, _ in self._SECTIONS:
+            cmds = section_cmds[label]
+            if not cmds:
+                continue
+            with formatter.section(f"{label}"):
+                formatter.write_dl(
+                    [(name, cmd.get_short_help_str(formatter.width)) for name, cmd in cmds]
+                )
+
+        if uncategorized:
+            with formatter.section("Other"):
+                formatter.write_dl(
+                    [(name, cmd.get_short_help_str(formatter.width)) for name, cmd in uncategorized]
+                )
+
+
+@click.group("dql", cls=_DqlGroup)
+def dql():
+    """Construct and run Diffbot Query Language (DQL) queries against the Knowledge Graph."""
+
+
+@dql.command("init")
+@click.option("--max-age", type=int, default=0, help="Skip the ontology re-download if the local copy is younger than N seconds (default 0 = always refresh).")
+@click.option("--keep-tmp", is_flag=True, help="Don't clear ~/.diffbot/tmp.")
+def init(max_age: int, keep_tmp: bool):
+    """Cache the Diffbot ontology and reset the tmp workspace."""
+    base = pathlib.Path.home() / ".diffbot"
+    base.mkdir(exist_ok=True)
+    tmp = base / "tmp"
+    if not keep_tmp:
+        if tmp.exists():
+            shutil.rmtree(tmp)
+        tmp.mkdir()
+
+    ont = ontology.ONTOLOGY_PATH
+    fresh = ont.exists() and max_age > 0 and (time.time() - ont.stat().st_mtime) < max_age
+    if fresh:
+        click.echo(f"ontology: cached ({int(time.time() - ont.stat().st_mtime)}s old)")
+    else:
+        get_client().dql_refresh_ontology(ont)
+        click.echo(f"ontology: refreshed -> {ont}")
+    ontology._CACHE.pop("data", None)
+
+    if resolve_token():
+        click.echo("credentials: token found")
+    else:
+        click.echo(
+            "credentials: MISSING -- no Diffbot API token found\n"
+            "  Set the DIFFBOT_API_TOKEN environment variable, or\n"
+            "  write 'DIFFBOT_API_TOKEN=YOUR_TOKEN' to ~/.diffbot/credentials",
+            err=True,
+        )
+        raise SystemExit(2)
+
+
+@dql.group("ontology")
+def ontology_group():
+    """Navigate the cached Diffbot ontology (types, fields, taxonomies, enums)."""
+
+
+@ontology_group.command("types")
+def ontology_types():
+    """List all entity type names."""
+    for n in ontology.list_types():
+        click.echo(n)
+
+
+@ontology_group.command("composites")
+def ontology_composites():
+    """List all composite type names."""
+    for n in ontology.list_composites():
+        click.echo(n)
+
+
+@ontology_group.command("enums")
+def ontology_enums():
+    """List all enum type names."""
+    for n in ontology.list_enums():
+        click.echo(n)
+
+
+@ontology_group.command("taxonomies")
+def ontology_taxonomies():
+    """List all taxonomy names."""
+    for n in ontology.list_taxonomies():
+        click.echo(n)
+
+
+@ontology_group.command("fields")
+@click.argument("type_name", metavar="TYPE")
+@click.argument("search", required=False, default=None)
+@click.option("--include-deprecated", is_flag=True)
+def ontology_fields(type_name: str, search, include_deprecated: bool):
+    """List fields of an entity type or composite (e.g. db dql ontology fields Organization)."""
+    fields = ontology.fields_for(type_name)
+    for name, meta in ontology.filter_fields(fields, search, include_deprecated=include_deprecated):
+        click.echo(ontology.format_field(name, meta))
+
+
+@ontology_group.command("taxonomy")
+@click.argument("name")
+@click.argument("search", required=False, default=None)
+def ontology_taxonomy(name: str, search):
+    """List values of a taxonomy."""
+    for v in ontology.taxonomy_values(name, search):
+        click.echo(v)
+
+
+@ontology_group.command("enum")
+@click.argument("name")
+def ontology_enum(name: str):
+    """List values of an enum."""
+    for v in ontology.enum_values(name):
+        click.echo(v)
+
+
+@ontology_group.command("search")
+@click.argument("term")
+def ontology_search(term: str):
+    """Generic fallback: search every 'name' field in the ontology by regex."""
+    for n in ontology.find_named(term):
+        click.echo(n)
+
+
+@dql.command("probe")
+@click.argument("queries", nargs=-1, required=True)
+@click.option("--workers", type=int, default=8, help="Maximum concurrent requests (default 8).")
+@click.option("--json", "as_json", is_flag=True, help="Emit results as a JSON array instead of a text table.")
+def probe(queries, workers: int, as_json: bool):
+    """Run multiple DQL queries in parallel and print hit counts (size=0) for each."""
+    reqs = [{"query": q, "size": 0} for q in queries]
+    results = get_client().dql_parallel(reqs, workers=workers)
+    rows = [
+        {"query": q, "hits": r.get("hits"), "results": r.get("results")}
+        for q, r in zip(queries, results)
+    ]
+    if as_json:
+        click.echo(json.dumps(rows, indent=2, ensure_ascii=False))
+        return
+    width = max(len(str(row["hits"])) for row in rows)
+    for row in rows:
+        click.echo(f"{str(row['hits']).rjust(width)}  {row['query']}")
+
+
+_DEFAULT_SPECS: dict[str, str] = {
+    "Organization": "name,Name;summary,Summary;nbEmployees,Employees;location.city.name,City;location.country.name,Country",
+    "Person":       "name,Name;$.employments[0].employer.name,Organization;$.employments[0].title,Title;location.city.name,City;location.country.name,Country",
+    "Article":      "date.str,Date;title,Title;author,Author;siteName,Site;language,Language",
+    "Product":      "name,Name;brand,Brand;category,Category;offerPrice,Price;summary,Summary",
+    "*":            "name,Name;description,Description",
+}
+
+# Extra fields appended to the stdout default spec for file exports (--out).
+# Exports can carry more columns than the terminal table; fill these in per type.
+# Format matches _DEFAULT_SPECS: "field.path,Display Name;..." (no leading ';').
+_EXPORT_EXTRA_SPECS: dict[str, str] = {
+    "Organization": "$.categories[*].name,Industries;revenue.value,Revenue;revenue.currency,Revenue Currency;foundingDate.str,Founding Date",
+    "Person":       "allNames,All Names;homepageUri,Website;linkedInUri,LinkedIn;$.employments[1:].employer.name,Employment History;",
+    "Article":      "authorUrl,Author URL;$.categories[*].name,Categories;tags.label,Tags;publisherCountry,Publisher Country;text,Text;html,HTML;",
+    "Product":      "",
+    "*":            "",
+}
+
+_TYPE_RE = re.compile(r"\btype:([A-Za-z]+)", re.IGNORECASE)
+
+
+def _type_for_query(query: str) -> str:
+    """Return the entity type named in the query, or '*' if none is found/known."""
+    m = _TYPE_RE.search(query)
+    if m and m.group(1) in _DEFAULT_SPECS:
+        return m.group(1)
+    return "*"
+
+
+def _spec_for_query(query: str) -> str:
+    """Return the default stdout exportspec for the entity type found in query."""
+    return _DEFAULT_SPECS[_type_for_query(query)]
+
+
+def _export_spec_for_query(query: str) -> str:
+    """Return the default file-export spec: the stdout default plus any per-type extras."""
+    t = _type_for_query(query)
+    base = _DEFAULT_SPECS[t]
+    extra = _EXPORT_EXTRA_SPECS.get(t, "").strip().strip(";")
+    return f"{base};{extra}" if extra else base
+
+
+@dql.command("export")
+@click.argument("query")
+@click.option("--out", default=None, help="Output file path. If omitted, results are printed to stdout as a table.")
+@click.option("--format", "fmt", type=click.Choice(["csv", "xls", "xlsx", "json"]), default="csv", help="Response format for --out (default csv).")
+@click.option("--spec", "exportspec", default=None, help="exportspec, e.g. 'name,Name;nbEmployees,Employees;location.city.name,City'")
+@click.option("--size", type=int, default=10, help="Page size (default 10).")
+@click.option("--from", "from_", type=int, default=0, help="Pagination offset.")
+@click.option("--filter", "filter_", default=None, help="Response field filter.")
+def export(query: str, out, fmt: str, exportspec, size: int, from_: int, filter_):
+    """Execute a DQL query and print to stdout or save to a file."""
+    if out is not None:
+        resolved_spec = exportspec or _export_spec_for_query(query)
+        body = get_client().dql(
+            query,
+            size=size,
+            from_=from_,
+            format=fmt,
+            exportspec=resolved_spec or None,
+            filter=filter_,
+            raw=True,
+        )
+        dest = pathlib.Path(out).expanduser()
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        dest.write_bytes(body)
+        click.echo(f"saved: {dest} ({len(body)} bytes)")
+        return
+
+    client = get_client()
+    resolved_spec = exportspec or _spec_for_query(query)
+
+    if sys.stdout.isatty():
+        with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+            progress.add_task(description=f"[Querying...] {query}", total=None)
+            hits = client.dql(query, size=0).get("hits", 0)
+            csv_bytes = client.dql(query, size=size, from_=from_, format="csv", exportspec=resolved_spec or None, filter=filter_, raw=True)
+    else:
+        hits = client.dql(query, size=0).get("hits", 0)
+        csv_bytes = client.dql(query, size=size, from_=from_, format="csv", exportspec=resolved_spec or None, filter=filter_, raw=True)
+
+    click.echo("")
+    click.echo(f"[Query] {query}")
+
+    reader = csv.reader(io.StringIO(csv_bytes.decode("utf-8")))
+    rows = list(reader)
+    if len(rows) < 2:
+        click.echo("(no results)")
+        return
+
+    headers, *data_rows = rows
+    table = Table(show_header=True, header_style="bold", show_lines=True)
+    for i, header in enumerate(headers):
+        if i < 3:
+            table.add_column(header, overflow="ellipsis", max_width=40, no_wrap=True)
+        else:
+            table.add_column(header, overflow="ellipsis", max_width=20, no_wrap=True)
+    for row in data_rows:
+        table.add_row(*row)
+
+    Console().print(table)
+
+    click.echo(f"{len(data_rows)} of {hits} records")
+    click.echo("")
+    click.echo("Options:")
+    click.echo("  --size N     return N records (default 10; -1 for all)")
+    click.echo("  --from N     skip the first N records (pagination offset)")
+    click.echo("")

diffbot/cli/entities.py

@@ -0,0 +1,155 @@
+"""CLI command: db entities — identify and resolve entities in text via the Diffbot NLP API."""
+
+import json
+import sys
+
+import click
+from rich.console import Console
+from rich.markup import escape
+from rich.progress import Progress, SpinnerColumn, TextColumn
+from rich.table import Table
+
+from diffbot import AuthError, APIError
+
+from ._common import get_client
+
+console = Console()
+
+_PALETTE = [
+    "cyan", "magenta", "yellow", "green", "blue", "red",
+    "bright_cyan", "bright_magenta", "bright_yellow", "bright_green",
+]
+
+
+def _sentiment_label(score) -> str:
+    if score is None:
+        return ""
+    if score > 0.1:
+        return f"[green]+{score:.2f}[/green]"
+    if score < -0.1:
+        return f"[red]{score:.2f}[/red]"
+    return f"[dim]{score:.2f}[/dim]"
+
+
+def _highlight_mentions(text: str, ents: list, palette: list) -> str:
+    spans = []
+    for i, e in enumerate(ents):
+        color = palette[i % len(palette)]
+        for m in e.get("mentions", []):
+            begin = m.get("beginOffset")
+            end = m.get("endOffset")
+            if begin is not None and end is not None:
+                spans.append((begin, end, color))
+
+    if not spans:
+        return escape(text)
+
+    spans.sort(key=lambda s: s[0])
+
+    parts = []
+    pos = 0
+    for begin, end, color in spans:
+        if begin < pos:
+            continue  # skip overlapping spans
+        parts.append(escape(text[pos:begin]))
+        parts.append(f"[bold underline {color}]{escape(text[begin:end])}[/bold underline {color}]")
+        pos = end
+    parts.append(escape(text[pos:]))
+    return "".join(parts)
+
+
+@click.command(name="entities")
+@click.argument("text", required=False)
+@click.option("-f", "--format", "fmt", type=click.Choice(["table", "json", "dql"]), default="table",
+              help="Output format: table (default), json (raw), dql (id:or(...) filter ready for db dql export)")
+@click.option("--lang", default="auto", help="Language code (default: auto)")
+def entities(text, fmt, lang):
+    """Identify and resolve entities and sentiment in text using the Diffbot NLP API.
+
+    TEXT can be passed as an argument or piped via stdin.
+
+    Entity IDs returned can be used in DQL for fast lookups:
+
+      db entities "..." -f dql | xargs db dql export
+    """
+    if text is None:
+        if sys.stdin.isatty():
+            raise click.UsageError("provide TEXT as an argument or pipe text via stdin")
+        text = sys.stdin.read().strip()
+        if not text:
+            raise click.UsageError("no text provided")
+
+    db = get_client()
+    try:
+        is_interactive = sys.stdout.isatty()
+        if is_interactive:
+            with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:
+                progress.add_task(description="Resolving entities...", total=None)
+                result = db.entities(text, lang=lang)
+        else:
+            result = db.entities(text, lang=lang)
+
+        if fmt == "json":
+            sys.stdout.write(json.dumps(result, indent=2))
+            sys.stdout.write("\n")
+            return
+
+        ents = result.get("entities", [])
+
+        if fmt == "dql":
+            ids = []
+            for e in ents:
+                raw = e.get("id") or e.get("diffbotUri") or ""
+                eid = raw.rstrip("/").rsplit("/", 1)[-1] if raw else ""
+                if eid:
+                    ids.append(eid)
+            if not ids:
+                return
+            sys.stdout.write('id:or(' + ",".join(f'"{i}"' for i in ids) + ')\n')
+            return
+
+        doc_sentiment = result.get("sentiment")
+        if doc_sentiment is not None:
+            console.print(f"Document sentiment: {_sentiment_label(doc_sentiment)}")
+            console.print()
+
+        if not ents:
+            console.print("[dim]No entities found.[/dim]")
+            return
+
+        table = Table(show_header=True, header_style="bold", show_lines=False)
+        table.add_column("Entity", overflow="fold")
+        table.add_column("Type", no_wrap=True)
+        table.add_column("Confidence", no_wrap=True, justify="right")
+        table.add_column("Salience", no_wrap=True, justify="right")
+        table.add_column("Sentiment", no_wrap=True, justify="right")
+        table.add_column("Diffbot ID", overflow="fold", style="dim")
+
+        for i, e in enumerate(ents):
+            color = _PALETTE[i % len(_PALETTE)]
+            name = e.get("name") or e.get("label") or ""
+            all_types = e.get("allTypes") or []
+            etype = all_types[0]["name"] if all_types else (e.get("type") or "")
+            confidence = e.get("confidence")
+            salience = e.get("salience")
+            sentiment = e.get("sentiment")
+            raw = e.get("id") or e.get("diffbotUri") or ""
+            eid = raw.rstrip("/").rsplit("/", 1)[-1] if raw else ""
+            conf_str = f"{confidence:.2f}" if confidence is not None else ""
+            sal_str = f"{salience:.2f}" if salience is not None else ""
+            table.add_row(
+                f"[{color}]● {name}[/{color}]",
+                etype, conf_str, sal_str, _sentiment_label(sentiment), eid,
+            )
+
+        console.print(_highlight_mentions(text, ents, _PALETTE))
+        console.print()
+        console.print(table)
+        console.print(f"[dim]{len(ents)} entit{'y' if len(ents) == 1 else 'ies'} found[/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()

diffbot/cli/ontology.py

@@ -0,0 +1,130 @@
+import json
+import pathlib
+import re
+from typing import Any, Dict, List, Optional
+
+ONTOLOGY_PATH = pathlib.Path.home() / ".diffbot" / "ontology.json"
+
+_CACHE: Dict[str, Any] = {}
+
+
+def load() -> Dict[str, Any]:
+    if "data" not in _CACHE:
+        if not ONTOLOGY_PATH.exists():
+            raise FileNotFoundError(
+                f"Ontology not found at {ONTOLOGY_PATH}. Run: db dql init"
+            )
+        _CACHE["data"] = json.loads(ONTOLOGY_PATH.read_text())
+    return _CACHE["data"]
+
+
+def list_types() -> List[str]:
+    return sorted(load().get("types", {}).keys())
+
+
+def list_composites() -> List[str]:
+    return sorted(load().get("composites", {}).keys())
+
+
+def list_enums() -> List[str]:
+    return sorted(load().get("enums", {}).keys())
+
+
+def list_taxonomies() -> List[str]:
+    return sorted(load().get("taxonomies", {}).keys())
+
+
+def _fields_of(container: Dict[str, Any], type_name: str) -> Dict[str, Any]:
+    entry = container.get(type_name)
+    if entry is None:
+        raise KeyError(f"Unknown name: {type_name}")
+    return entry.get("fields", {})
+
+
+def fields_for(type_name: str) -> Dict[str, Any]:
+    data = load()
+    types = data.get("types", {})
+    composites = data.get("composites", {})
+    if type_name in types:
+        return _fields_of(types, type_name)
+    if type_name in composites:
+        return _fields_of(composites, type_name)
+    raise KeyError(f"{type_name} is not a known entity type or composite")
+
+
+def format_field(name: str, meta: Dict[str, Any]) -> str:
+    t = meta.get("type", "?")
+    if t == "LinkedEntity":
+        le = meta.get("leType") or []
+        if le:
+            t = f"LinkedEntity ({le[0]})"
+    flags = []
+    if meta.get("isList"):
+        flags.append("isList")
+    if meta.get("isComposite"):
+        flags.append("isComposite")
+    if meta.get("isEnum"):
+        flags.append("isEnum")
+    if meta.get("isDeprecated"):
+        flags.append("DEPRECATED")
+    suffix = "".join(f" [{f}]" for f in flags)
+    return f"{name}: [{t}]{suffix}"
+
+
+def filter_fields(fields: Dict[str, Any], search: Optional[str], include_deprecated: bool = False) -> List[tuple]:
+    pattern = re.compile(search, re.IGNORECASE) if search else None
+    out = []
+    for name, meta in fields.items():
+        if not include_deprecated and meta.get("isDeprecated"):
+            continue
+        if pattern and not pattern.search(name):
+            continue
+        out.append((name, meta))
+    return out
+
+
+def taxonomy_values(name: str, search: Optional[str] = None) -> List[str]:
+    data = load()
+    tax = data.get("taxonomies", {}).get(name)
+    if tax is None:
+        raise KeyError(f"Unknown taxonomy: {name}")
+    pattern = re.compile(search, re.IGNORECASE) if search else None
+    out: List[str] = []
+
+    def walk(node: Dict[str, Any]) -> None:
+        n = node.get("name")
+        if n and (pattern is None or pattern.search(n)):
+            out.append(n)
+        for child in node.get("children", []) or []:
+            walk(child)
+
+    for cat in tax.get("categories", []) or []:
+        walk(cat)
+    return out
+
+
+def enum_values(name: str) -> List[str]:
+    data = load()
+    enum = data.get("enums", {}).get(name)
+    if enum is None:
+        raise KeyError(f"Unknown enum: {name}")
+    return list(enum.get("values", []))
+
+
+def find_named(search: str) -> List[str]:
+    pattern = re.compile(search, re.IGNORECASE)
+    found = set()
+
+    def walk(node: Any) -> None:
+        if isinstance(node, dict):
+            n = node.get("name")
+            if isinstance(n, str) and pattern.search(n):
+                found.add(n)
+            for v in node.values():
+                walk(v)
+        elif isinstance(node, list):
+            for v in node:
+                walk(v)
+
+    walk(load())
+    return sorted(found)