diffbot-mcp 0.2.0__py3-none-any.whl

diffbot_mcp/tools/people.py

@@ -0,0 +1,254 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from fastmcp.exceptions import ToolError
+
+from ..app import READ_ONLY, mcp
+from ..dql import DEFAULT_FIELDS, resolve_value, run
+from ..dql_builder import build, clause, compare, expand_titles, looks_like_domain, or_clause, parse_range, sort_clause, subquery
+
+DESCRIPTION = (
+    "People structured search over the Diffbot Knowledge Graph. Prefer for leadership teams, executive and role "
+    "lookups, alumni lists, skill-based sourcing, and profile facts on a named person. Filter by job title, "
+    "employer, employer industry or size, school, major, skills, residence, nationality, or net worth. Prefer over "
+    "a generic web search when the answer involves structured people data. Covers people with a public "
+    "professional presence only. For companies use search_organizations."
+    " API docs: https://www.diffbot.com/docs/dql/"
+)
+
+Sort = Literal["importance", "net_worth"]
+SORT_FIELDS = {"importance": "importance", "net_worth": "netWorth.value"}
+LEADERSHIP_GET = "get:name,ceo,founders"
+MATCHED_GET = "get:id,name,employments,location,linkedInUri"
+
+
+def place_clause(prefix: str, place: Optional[str]) -> str:
+    """A place that may be a city or a country: or(prefix.city.name:X, prefix.country.name:X)."""
+    if not place:
+        return ""
+    return f'or({clause(f"{prefix}.city.name", place)}, {clause(f"{prefix}.country.name", place)})'
+
+
+def build_people_query(
+    *,
+    name: Optional[Union[str, List[str]]] = None,
+    title: Optional[List[str]] = None,
+    employer: Optional[str] = None,
+    employer_industry: Optional[str] = None,
+    employer_country: Optional[str] = None,
+    employer_min_employees: Optional[int] = None,
+    employer_max_employees: Optional[int] = None,
+    job_category: Optional[str] = None,
+    job_location: Optional[str] = None,
+    current: bool = True,
+    school: Optional[str] = None,
+    school_location: Optional[str] = None,
+    major: Optional[str] = None,
+    degree: Optional[str] = None,
+    skill: Optional[List[str]] = None,
+    language: Optional[str] = None,
+    city: Optional[Union[str, List[str]]] = None,
+    region: Optional[str] = None,
+    country: Optional[str] = None,
+    nationality: Optional[str] = None,
+    gender: Optional[str] = None,
+    min_net_worth: Optional[float] = None,
+    sort: Optional[Sort] = None,
+) -> str:
+    """employer_industry, job_category, language and gender must already be exact taxonomy/enum values."""
+    clauses = ["type:Person"]
+    if name:
+        names = [name] if isinstance(name, str) else list(name)
+        clauses.append(or_clause("name", names, strict=True))
+
+    # every condition about one job goes inside one {}; written separately they match
+    # different jobs and inflate the count roughly nine times
+    employment = []
+    if title:
+        employment.append(or_clause("title", expand_titles(title)))
+    if employer:
+        employment.append(clause("employer.name", employer))
+    if employer_industry:
+        employment.append(clause("employer.categories.name", employer_industry))
+    if employer_country:
+        employment.append(clause("employer.location.country.name", employer_country))
+    if employer_min_employees is not None:
+        employment.append(compare("employer.nbEmployees", ">=", employer_min_employees))
+    if employer_max_employees is not None:
+        employment.append(compare("employer.nbEmployees", "<=", employer_max_employees))
+    if job_category:
+        employment.append(clause("categories.name", job_category))
+    if job_location:
+        employment.append(place_clause("location", job_location))
+    if employment and current:
+        employment.append("isCurrent:true")
+    clauses.append(subquery("employments", employment))
+
+    education = []
+    if school:
+        education.append(clause("institution.name", school))
+    if major:
+        education.append(clause("major.name", major))
+    if degree:
+        education.append(clause("degree.name", degree))
+    if school_location:
+        education.append(place_clause("location", school_location))
+    clauses.append(subquery("educations", education))
+
+    if skill:
+        clauses.append(or_clause("skills.name", skill))
+    if language:
+        clauses.append(clause("languages", language))
+    for level, value in (("city", city), ("region", region), ("country", country)):
+        if value:
+            values = [value] if isinstance(value, str) else list(value)
+            clauses.append(or_clause(f"location.{level}.name", values))
+    if nationality:
+        clauses.append(clause("nationalities.name", nationality))
+    if gender:
+        clauses.append(clause("gender", gender))
+    if min_net_worth is not None:
+        clauses.append(compare("netWorth.value", ">=", min_net_worth))
+
+    if len([c for c in clauses if c]) == 1:
+        raise ToolError("search_people needs at least one filter.")
+    if sort:
+        clauses.append(sort_clause(SORT_FIELDS[sort], descending=True))
+    return build(clauses)
+
+
+def build_leadership_query(company: str) -> str:
+    # the domain plus the default ranking is right for apple.com, tesla.com, openai.com and
+    # friends, where strict:name misses (Apple Inc.) and a sort surfaces the wrong record
+    anchor = clause("homepageUri", company) if looks_like_domain(company) else clause("name", company)
+    return build(["type:Organization", anchor, LEADERSHIP_GET])
+
+
+def pick_employment(employments: List[dict], employer: Optional[str], titles: List[str], current: bool) -> dict:
+    """The employment that best matches the filters, so rendered columns reflect the matched job rather than the primary one."""
+    def score(job: dict) -> int:
+        points = 0
+        employer_name = ((job.get("employer") or {}).get("name") or "").lower()
+        if employer and employer.lower() in employer_name:
+            points += 4
+        job_title = (job.get("title") or "").lower()
+        if titles and any(t.lower() in job_title for t in titles):
+            points += 2
+        if current and job.get("isCurrent"):
+            points += 1
+        return points
+
+    return max(employments, key=score) if employments else {}
+
+
+def render_person(entity: dict, employer: Optional[str], titles: List[str], current: bool) -> dict:
+    job = pick_employment(entity.get("employments") or [], employer, titles, current)
+    location = entity.get("location") or {}
+    row = {
+        "Name": entity.get("name"),
+        "Title": job.get("title"),
+        "Employer": (job.get("employer") or {}).get("name"),
+        "City": (location.get("city") or {}).get("name"),
+        "Country": (location.get("country") or {}).get("name"),
+        "LinkedIn": entity.get("linkedInUri"),
+        "Id": entity.get("id"),
+    }
+    return {key: value for key, value in row.items() if value}
+
+
+@mcp.tool(name="search_people", title="Search People", description=DESCRIPTION, annotations=READ_ONLY)
+async def search_people(
+    name: Annotated[Optional[Union[str, List[str]]], "Person name(s), exact"] = None,
+    title: Annotated[Optional[List[str]], "Job titles (OR); C-suite expanded"] = None,
+    employer: Annotated[Optional[str], "Employer name"] = None,
+    employer_industry: Annotated[Optional[str], "Employer's industry; free text ok"] = None,
+    employer_employees: Annotated[Optional[str], "Employer headcount range, e.g. '<=50'"] = None,
+    employer_country: Annotated[Optional[str], "Employer HQ country"] = None,
+    job_category: Annotated[Optional[str], "Occupation group, e.g. 'Business', 'Religion'"] = None,
+    job_location: Annotated[Optional[str], "City or country where the job is"] = None,
+    current: Annotated[bool, "False: include past jobs"] = True,
+    school: Annotated[Optional[str], "Institution name"] = None,
+    school_location: Annotated[Optional[str], "City or country of the school"] = None,
+    major: Annotated[Optional[str], "Field of study"] = None,
+    degree: Annotated[Optional[str], "Degree, e.g. 'MBA'"] = None,
+    skill: Annotated[Optional[List[str]], "Skills (OR)"] = None,
+    language: Annotated[Optional[str], "Language spoken"] = None,
+    city: Annotated[Optional[Union[str, List[str]]], "Residence city or cities"] = None,
+    region: Annotated[Optional[str], "Residence state/region"] = None,
+    country: Annotated[Optional[str], "Residence country"] = None,
+    nationality: Annotated[Optional[str], "Nationality as country name"] = None,
+    gender: Annotated[Optional[str], "Gender"] = None,
+    net_worth: Annotated[Optional[str], "Minimum net worth, e.g. '1B'"] = None,
+    sort: Annotated[Optional[Sort], "Only if ordering is the question"] = None,
+    leadership_of: Annotated[Optional[str], "Company domain or name: its CEO and founders"] = None,
+    size: Annotated[int, "Max rows"] = 10,
+    offset: Annotated[int, "Rows to skip"] = 0,
+) -> dict:
+    if leadership_of:
+        result = await run(build_leadership_query(leadership_of), size=1, format="json")
+        records = result.get("results") or []
+        if not records:
+            return {"dql": result["dql"], "hits": result.get("hits"), "company": None, "ceo": None, "founders": [], "notes": ["No organization matched. Prefer the homepage domain over the name."]}
+        entity = records[0]
+        return {
+            "dql": result["dql"],
+            "hits": result.get("hits"),
+            "company": entity.get("name"),
+            "ceo": (entity.get("ceo") or {}).get("name"),
+            "founders": [f.get("name") for f in entity.get("founders") or [] if f.get("name")],
+        }
+
+    notes: List[str] = []
+    if employer_industry:
+        resolved = await resolve_value("OrganizationCategory", employer_industry)
+        if resolved.value is None:
+            if resolved.candidates:
+                return {"dql": None, "candidates": {"employer_industry": resolved.candidates}, "notes": [resolved.note]}
+            raise ToolError(resolved.note)
+        employer_industry = resolved.value
+        if resolved.note:
+            notes.append(resolved.note)
+    if job_category:
+        resolved = await resolve_value("EmploymentCategory", job_category)
+        if resolved.value is None:
+            if resolved.candidates:
+                return {"dql": None, "candidates": {"job_category": resolved.candidates}, "notes": [resolved.note]}
+            raise ToolError(resolved.note)
+        job_category = resolved.value
+        if resolved.note:
+            notes.append(resolved.note)
+    for label, kind_name in (("gender", "Gender"), ("language", "Language")):
+        value = gender if label == "gender" else language
+        if value:
+            resolved = await resolve_value(kind_name, value, kind="enum")
+            if resolved.value is None:
+                if resolved.candidates:
+                    return {"dql": None, "candidates": {label: resolved.candidates}, "notes": [resolved.note]}
+                raise ToolError(resolved.note)
+            if label == "gender":
+                gender = resolved.value
+            else:
+                language = resolved.value
+    employer_min, employer_max = parse_range(employer_employees, "employer_employees", integer=True)
+    min_net_worth, _ = parse_range(net_worth, "net_worth")
+
+    query = build_people_query(
+        name=name, title=title, employer=employer, employer_industry=employer_industry, employer_country=employer_country,
+        employer_min_employees=employer_min, employer_max_employees=employer_max, job_category=job_category,
+        job_location=job_location, current=current, school=school, school_location=school_location, major=major,
+        degree=degree, skill=skill, language=language, city=city, region=region, country=country,
+        nationality=nationality, gender=gender, min_net_worth=min_net_worth, sort=sort,
+    )
+
+    # with an employer or title filter the rendered job must be the one that matched, not
+    # the primary one, so the records are read as json and the matching employment picked
+    if employer or title:
+        result = await run(f"{query} {MATCHED_GET}", size=size, offset=offset, format="json", bare_type_check=True)
+        titles = expand_titles(title or [])
+        result["results"] = [render_person(entity, employer, titles, current) for entity in result.get("results") or []]
+        result["dql"] = query
+    else:
+        result = await run(query, size=size, offset=offset, fields=DEFAULT_FIELDS["Person"], bare_type_check=True)
+
+    if notes:
+        result["notes"] = result.get("notes", []) + notes
+    return result

diffbot_mcp/tools/places.py

@@ -0,0 +1,109 @@
+from typing import Annotated, List, Literal, Optional
+
+from fastmcp.exceptions import ToolError
+
+from ..app import READ_ONLY, mcp
+from ..dql import PLACE_FIELDS, run
+from ..dql_builder import build, clause, compare, name_variants, near_clause, or_clause, sort_clause, parse_range, split_near
+
+DESCRIPTION = (
+    "Cities, counties, states, provinces, countries, and points of interest in the Diffbot Knowledge Graph by "
+    "name, population, containing place, or proximity. Prefer for facts on places, rather than recalling from "
+    "memory."
+    " API docs: https://www.diffbot.com/docs/dql/"
+)
+
+Kind = Literal["city", "subregion", "region", "country", "administrative_area", "poi", "any"]
+KIND_CLAUSES = {
+    "city": "type:City",
+    "subregion": "type:Subregion",
+    "region": "type:Region",
+    "country": "type:Country",
+    "administrative_area": "type:AdministrativeArea",
+    "poi": 'type:Place not(types:"AdministrativeArea")',
+    "any": "type:Place",
+}
+Sort = Literal["population", "importance"]
+
+
+def build_places_query(
+    *,
+    kind: Kind = "any",
+    name: Optional[str] = None,
+    descriptors: Optional[str] = None,
+    country: Optional[str] = None,
+    region: Optional[str] = None,
+    subregion: Optional[str] = None,
+    within: Optional[str] = None,
+    continent: Optional[str] = None,
+    near: Optional[str] = None,
+    radius: Optional[str] = None,
+    min_population: Optional[int] = None,
+    max_population: Optional[int] = None,
+    sort: Optional[Sort] = None,
+    facet: Optional[Literal["placeType"]] = None,
+) -> tuple:
+    """Return (dql, notes)."""
+    clauses = [KIND_CLAUSES[kind]]
+    notes: List[str] = []
+
+    if name:
+        # diacritics are not normalized and coverage of each spelling is inconsistent
+        clauses.append(or_clause("name", name_variants(name), strict=True))
+    if descriptors:
+        clauses.append(clause("descriptors", descriptors))
+    for level, value in (("country", country), ("region", region), ("subregion", subregion)):
+        if value:
+            clauses.append(clause(f"location.{level}.name", value))
+    if within:
+        clauses.append(clause("isPartOf.name", within))
+    if continent:
+        clauses.append(clause("description", f"in {continent}"))
+        notes.append(f"There is no continent field; '{continent}' was matched against the prose description, which is over-inclusive and double-counts duplicate records. Filter the list before presenting it and say it was text-matched.")
+    if near:
+        clauses.append(near_clause(near, radius))
+    if min_population is not None:
+        clauses.append(compare("population", ">=", min_population))
+    if max_population is not None:
+        clauses.append(compare("population", "<=", max_population))
+
+    if len(clauses) == 1 and kind in ("any", "administrative_area", "poi"):
+        raise ToolError("search_places needs a name, containing place, or other filter for this kind.")
+
+    if facet:
+        clauses.append(f"facet:{facet}")
+    elif sort:
+        clauses.append(sort_clause(sort, descending=True))
+    return build(clauses), notes
+
+
+@mcp.tool(name="search_places", title="Search Places", description=DESCRIPTION, annotations=READ_ONLY)
+async def search_places(
+    kind: Annotated[Kind, "country | region | subregion | city | administrative_area | poi | any"] = "any",
+    name: Annotated[Optional[str], "Place name, exact"] = None,
+    descriptors: Annotated[Optional[str], "Free text, e.g. 'national park' (poi)"] = None,
+    country: Annotated[Optional[str], "Containing country"] = None,
+    region: Annotated[Optional[str], "Containing state/region"] = None,
+    subregion: Annotated[Optional[str], "Containing county/district"] = None,
+    within: Annotated[Optional[str], "Containing place, any level"] = None,
+    continent: Annotated[Optional[str], "Continent (text match, loose)"] = None,
+    near: Annotated[Optional[str], "Place and radius, e.g. 'Paris 50km'"] = None,
+    population: Annotated[Optional[str], "Range, e.g. '1M+', '10000-50000'"] = None,
+    sort: Annotated[Optional[Sort], "Only if ordering is the question"] = None,
+    facet: Annotated[Optional[Literal["placeType"]], "Bucket by place type"] = None,
+    size: Annotated[int, "Max rows"] = 10,
+    offset: Annotated[int, "Rows to skip"] = 0,
+) -> dict:
+    near, radius = split_near(near)
+    min_population, max_population = parse_range(population, "population", integer=True)
+    query, notes = build_places_query(
+        kind=kind, name=name, descriptors=descriptors, country=country, region=region, subregion=subregion,
+        within=within, continent=continent, near=near, radius=radius, min_population=min_population,
+        max_population=max_population, sort=sort, facet=facet,
+    )
+    result = await run(query, size=size, offset=offset, fields=PLACE_FIELDS, bare_type_check=True)
+    if result.get("hits") == 0 and name:
+        notes.append("Zero hits for both spellings of the name. Try kind 'any', or match on country/region without the name.")
+    if notes:
+        result["notes"] = result.get("notes", []) + notes
+    return result

diffbot_mcp/tools/resolve_entities.py

@@ -0,0 +1,52 @@
+from typing import Annotated, Optional
+
+from ..api import request, resolve_token
+from ..app import READ_ONLY, mcp
+
+NLP_ENDPOINT = "https://nl.diffbot.com/v1/"
+
+DESCRIPTION = (
+    "Named entity recognition and linking. Finds people, organizations, places, and products mentioned in text. "
+    "Each resolves to a Diffbot Knowledge Graph entity with confidence, salience, and sentiment. Prefer for "
+    "post-processing plain text for entity mentions. The response includes dql_filter, an id:or(...) clause for "
+    "dql to fetch the full records. Entities without an id were recognized but not matched."
+    " API docs: https://www.diffbot.com/docs/natural-language/"
+)
+
+
+def entity_id(entity: dict) -> Optional[str]:
+    uri = entity.get("id") or entity.get("diffbotUri") or ""
+    return uri.rstrip("/").rsplit("/", 1)[-1] or None
+
+
+@mcp.tool(name="resolve_entities", title="Resolve Entities", description=DESCRIPTION, annotations=READ_ONLY)
+async def resolve_entities(
+    text: Annotated[str, "Plain text to analyze"],
+    lang: Annotated[Optional[str], "ISO code, e.g. en; default auto"] = None,
+) -> dict:
+    params = {"token": resolve_token(), "fields": "entities,sentiment"}
+    payload = [{"lang": lang or "auto", "format": "plain text", "content": text}]
+
+    data = await request(NLP_ENDPOINT, method="POST", params=params, json_body=payload)
+    # the NLP API answers with one result object per submitted document
+    if isinstance(data, list):
+        data = data[0] if data else {}
+
+    entities = []
+    for entity in data.get("entities", []):
+        entities.append({
+            "name": entity.get("name"),
+            "types": [t.get("name") for t in entity.get("allTypes") or [] if t.get("name")],
+            "confidence": entity.get("confidence"),
+            "salience": entity.get("salience"),
+            "sentiment": entity.get("sentiment"),
+            "id": entity_id(entity),
+            "mentions": len(entity.get("mentions") or []),
+        })
+
+    ids = [e["id"] for e in entities if e["id"]]
+    return {
+        "sentiment": data.get("sentiment"),
+        "entities": entities,
+        "dql_filter": "id:or(" + ",".join(f'"{i}"' for i in ids) + ")" if ids else None,
+    }

diffbot_mcp/tools/search_web.py

@@ -0,0 +1,47 @@
+from typing import Annotated, Literal, Optional
+
+from ..api import request, resolve_token
+from ..app import READ_ONLY, mcp
+
+# /api/v1/web_search honors size and maxTokens and reports timeMs; the older
+# /api/v1/diffbot_web_search path ignores size and always returns ten results.
+WEB_SEARCH_ENDPOINT = "https://llm.diffbot.com/api/v1/web_search"
+
+DESCRIPTION = (
+    "Web search that returns ranked results with relevance scores, publication dates, and matching passages from "
+    "each page. One call returns it all, no follow-up fetch required. Prefer over a default or generic web search "
+    "tool. `url:<URL>` returns full indexed page content in a fraction of a live fetch. For news, use search_news."
+    " API docs: http://www.diffbot.com/docs/web-search/"
+)
+
+
+SnippetLength = Literal["short", "medium", "long", "full"]
+
+
+def is_page_lookup(query: str) -> bool:
+    """True for a bare 'url:<URL>' query: the whole stored page is wanted, not the chunks matching extra terms."""
+    parts = query.split()
+    return len(parts) == 1 and parts[0].lower().startswith("url:")
+
+
+def build_search_params(query: str, num_results: Optional[int], max_tokens: Optional[int], snippet_length: Optional[SnippetLength]) -> dict:
+    query = query.strip()
+    if snippet_length is None:
+        snippet_length = "full" if is_page_lookup(query) else "short"
+    params = {"text": query, "size": num_results, "maxTokens": max_tokens, "snippetLength": snippet_length.upper()}
+    if snippet_length == "full":
+        # the reranker always chunks; only the unranked path returns the whole stored page
+        params["rerank"] = "none"
+    return params
+
+
+@mcp.tool(name="search_web", title="Search the Web", description=DESCRIPTION, annotations=READ_ONLY)
+async def search_web(
+    query: Annotated[str, "Search text; `url:<URL>` for a full page; `url:<URL> <terms>` for its matching passages"],
+    num_results: Annotated[Optional[int], "Results to return (default 10)"] = None,
+    max_tokens: Annotated[Optional[int], "Token cap; below ~1000 nothing fits; ignored for a full page"] = None,
+    snippet_length: Annotated[Optional[SnippetLength], "Passage size per result; default short (~1-1.5 KB), full for a bare `url:` (the whole page, terms ignored)"] = None,
+) -> dict:
+    params = build_search_params(query, num_results, max_tokens, snippet_length)
+    headers = {"Authorization": f"Bearer {resolve_token()}"}
+    return await request(WEB_SEARCH_ENDPOINT, params=params, headers=headers)