langchain-diffbot 0.1.0__py3-none-any.whl

langchain_diffbot/__init__.py

@@ -0,0 +1,39 @@
+"""LangChain integration for Diffbot.
+
+Thin layer over the official `diffbot-python` SDK. Every public class accepts
+either a `diffbot_api_token` (or `DIFFBOT_API_TOKEN` env var) or a pre-built
+`diffbot.Diffbot` / `diffbot.DiffbotAsync` client via the `client` /
+`async_client` fields — anything the SDK can do, you can do via these classes.
+"""
+
+from langchain_diffbot.chat_models import ChatDiffbot
+from langchain_diffbot.document_loaders import (
+    DiffbotCrawlLoader,
+    DiffbotExtractLoader,
+)
+from langchain_diffbot.retrievers import (
+    DiffbotKnowledgeGraphRetriever,
+    DiffbotWebSearchRetriever,
+)
+from langchain_diffbot.tools import (
+    DiffbotDQLProbeTool,
+    DiffbotEntitiesTool,
+    DiffbotExtractTool,
+    DiffbotKnowledgeGraphTool,
+    DiffbotOntologyTool,
+    DiffbotWebSearchTool,
+)
+
+__all__ = [
+    "ChatDiffbot",
+    "DiffbotCrawlLoader",
+    "DiffbotDQLProbeTool",
+    "DiffbotEntitiesTool",
+    "DiffbotExtractLoader",
+    "DiffbotExtractTool",
+    "DiffbotKnowledgeGraphRetriever",
+    "DiffbotKnowledgeGraphTool",
+    "DiffbotOntologyTool",
+    "DiffbotWebSearchRetriever",
+    "DiffbotWebSearchTool",
+]

langchain_diffbot/_base.py

@@ -0,0 +1,110 @@
+"""Shared base for langchain-diffbot components.
+
+Every public class in this package inherits from `_BaseDiffbotComponent`. The
+mixin holds the token / timeout / optional pre-built SDK clients and exposes
+two context managers (`_sync_db`, `_async_db`) that the components use to
+acquire a `diffbot.Diffbot` / `diffbot.DiffbotAsync` for a single call.
+
+Bring-your-own-client: if the user supplies `client=...` or
+`async_client=...`, we use it as-is and **do not close it** — the user owns
+the lifecycle. Otherwise we construct a fresh SDK client per call and close
+it on exit (same per-call lifecycle as the previous hand-rolled httpx
+wrapper).
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
+
+from diffbot import Diffbot, DiffbotAsync
+from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
+
+
+class _BaseDiffbotComponent(BaseModel):
+    """Mixin holding token, timeout, and optional pre-built SDK clients.
+
+    Concrete classes inherit from this *and* a LangChain base
+    (`BaseRetriever`, `BaseTool`, `BaseDocumentLoader`, `BaseChatModel`).
+    Both are Pydantic models, so their fields merge.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    diffbot_api_token: SecretStr | None = Field(default=None)
+    """Diffbot API token. Falls back to `DIFFBOT_API_TOKEN`.
+
+    Not required when both `client` and `async_client` are supplied.
+    """
+
+    timeout: float = 30.0
+    """HTTP timeout (seconds) for SDK clients we construct ourselves.
+
+    Ignored when `client` / `async_client` are supplied.
+    """
+
+    client: Diffbot | None = Field(default=None, exclude=True, repr=False)
+    """Optional pre-built sync SDK client.
+
+    If set, we use it as-is and do not close it.
+    """
+
+    async_client: DiffbotAsync | None = Field(default=None, exclude=True, repr=False)
+    """Optional pre-built async SDK client.
+
+    If set, we use it as-is and do not close it.
+    """
+
+    @model_validator(mode="after")
+    def _resolve_token(self) -> _BaseDiffbotComponent:
+        # If the user gave us a client (for either side), we can't be sure
+        # they'll use the other side — but token resolution shouldn't block
+        # construction in that case. Defer the missing-token error to call time.
+        if self.client is not None or self.async_client is not None:
+            return self
+        if (
+            self.diffbot_api_token is None
+            or not self.diffbot_api_token.get_secret_value()
+        ):
+            env_token = os.environ.get("DIFFBOT_API_TOKEN", "")
+            if not env_token:
+                msg = (
+                    "A Diffbot API token is required. Pass `diffbot_api_token=...`, "
+                    "set the `DIFFBOT_API_TOKEN` environment variable, or supply a "
+                    "pre-built `client` / `async_client`."
+                )
+                raise ValueError(msg)
+            self.diffbot_api_token = SecretStr(env_token)
+        return self
+
+    def _token(self) -> str:
+        if (
+            self.diffbot_api_token is None
+            or not self.diffbot_api_token.get_secret_value()
+        ):
+            msg = (
+                "A Diffbot API token is required for this call. Pass "
+                "`diffbot_api_token=...`, set `DIFFBOT_API_TOKEN`, or supply a "
+                "pre-built client."
+            )
+            raise ValueError(msg)
+        return self.diffbot_api_token.get_secret_value()
+
+    @contextmanager
+    def _sync_db(self) -> Iterator[Diffbot]:
+        """Yield a `Diffbot` for one call. Closes only clients we constructed."""
+        if self.client is not None:
+            yield self.client
+            return
+        with Diffbot(token=self._token(), timeout=self.timeout) as db:
+            yield db
+
+    @asynccontextmanager
+    async def _async_db(self) -> AsyncIterator[DiffbotAsync]:
+        """Yield a `DiffbotAsync` for one call. Closes only clients we constructed."""
+        if self.async_client is not None:
+            yield self.async_client
+            return
+        async with DiffbotAsync(token=self._token(), timeout=self.timeout) as db:
+            yield db

langchain_diffbot/chat_models.py

@@ -0,0 +1,122 @@
+"""ChatDiffbot — LangChain chat model wrapping Diffbot's LLM RAG `ask` endpoint."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from typing import Any
+
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
+
+from langchain_diffbot._base import _BaseDiffbotComponent
+
+
+def _to_diffbot_messages(messages: list[BaseMessage]) -> list[dict[str, str]]:
+    out: list[dict[str, str]] = []
+    for m in messages:
+        if isinstance(m, HumanMessage):
+            role = "user"
+        elif isinstance(m, AIMessage):
+            role = "assistant"
+        elif isinstance(m, SystemMessage):
+            role = "system"
+        else:
+            # Fall back on the message's `type` attribute for anything exotic.
+            role = m.type if isinstance(m.type, str) else "user"
+        content = m.content if isinstance(m.content, str) else str(m.content)
+        out.append({"role": role, "content": content})
+    return out
+
+
+class ChatDiffbot(_BaseDiffbotComponent, BaseChatModel):
+    """Chat model backed by Diffbot's LLM RAG (`ask`) endpoint.
+
+    The SDK streams tokens natively, so this class implements `_stream` and
+    `_astream`. `_generate` / `_agenerate` aggregate the stream into a single
+    `ChatGeneration`.
+
+    Example:
+        ```python
+        from langchain_diffbot import ChatDiffbot
+
+        llm = ChatDiffbot()
+        llm.invoke("What's the capital of France?")
+        ```
+    """
+
+    @property
+    def _llm_type(self) -> str:
+        return "diffbot"
+
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        payload = _to_diffbot_messages(messages)
+        with self._sync_db() as db:
+            for chunk in db.ask(payload):
+                if run_manager is not None:
+                    run_manager.on_llm_new_token(chunk)
+                yield ChatGenerationChunk(message=AIMessageChunk(content=chunk))
+
+    async def _astream(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        payload = _to_diffbot_messages(messages)
+        async with self._async_db() as db:
+            async for chunk in db.ask(payload):
+                if run_manager is not None:
+                    await run_manager.on_llm_new_token(chunk)
+                yield ChatGenerationChunk(message=AIMessageChunk(content=chunk))
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        parts: list[str] = []
+        for chunk in self._stream(
+            messages, stop=stop, run_manager=run_manager, **kwargs
+        ):
+            content = chunk.message.content
+            parts.append(content if isinstance(content, str) else str(content))
+        return ChatResult(
+            generations=[ChatGeneration(message=AIMessage(content="".join(parts)))]
+        )
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        stop: list[str] | None = None,
+        run_manager: AsyncCallbackManagerForLLMRun | None = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        parts: list[str] = []
+        async for chunk in self._astream(
+            messages, stop=stop, run_manager=run_manager, **kwargs
+        ):
+            content = chunk.message.content
+            parts.append(content if isinstance(content, str) else str(content))
+        return ChatResult(
+            generations=[ChatGeneration(message=AIMessage(content="".join(parts)))]
+        )

langchain_diffbot/document_loaders.py

@@ -0,0 +1,144 @@
+"""Diffbot document loaders — batch ingestion via Extract and Crawl."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable, Iterator
+from typing import Any
+
+from diffbot.crawl import CrawlEvent, CrawlEventType
+from langchain_core.document_loaders import BaseLoader
+from langchain_core.documents import Document
+from pydantic import Field
+
+from langchain_diffbot._base import _BaseDiffbotComponent
+
+ExtractDocumentMapper = Callable[[str, dict[str, Any]], Document]
+"""`(url, raw_response) -> Document`."""
+
+CrawlEventMapper = Callable[[CrawlEvent], Document | None]
+"""`(event) -> Document | None`. Return None to skip the event."""
+
+
+def _default_extract_mapper(url: str, raw: dict[str, Any]) -> Document:
+    objects = raw.get("objects") or []
+    first = objects[0] if objects else {}
+    page_content = first.get("text") or raw.get("markdown") or ""
+    metadata = {
+        "url": url,
+        "title": first.get("title") or raw.get("title"),
+        "pageUrl": first.get("pageUrl") or raw.get("url"),
+        "resolvedPageUrl": first.get("resolvedPageUrl"),
+        "type": first.get("type") or raw.get("type"),
+    }
+    return Document(page_content=page_content, metadata=metadata)
+
+
+def _default_crawl_mapper(event: CrawlEvent) -> Document | None:
+    # Non-URL events (e.g. JOB_CREATED) are skipped by default. The Document's
+    # page_content is the URL itself — the crawl SDK only surfaces URLs, not
+    # page contents. Chain with `DiffbotExtractLoader` to fetch content.
+    if event.event_type != CrawlEventType.URL_PROCESSED:
+        return None
+    url = event.details.get("url", "")
+    return Document(
+        page_content=url,
+        metadata={
+            "url": url,
+            "status": event.details.get("status"),
+            "crawl_timestamp": event.timestamp,
+        },
+    )
+
+
+class DiffbotExtractLoader(_BaseDiffbotComponent, BaseLoader):
+    """Loader that calls `extract` on each of `urls` and yields a `Document`.
+
+    Example:
+        ```python
+        from langchain_diffbot import DiffbotExtractLoader
+
+        docs = DiffbotExtractLoader(
+            urls=["https://example.com", "https://diffbot.com"],
+        ).load()
+        ```
+    """
+
+    urls: list[str]
+    """URLs to extract."""
+
+    api: str = "analyze"
+    """Diffbot extract API. Defaults to `analyze`."""
+
+    fmt: str = "markdown"
+    """Output format. `markdown` uses Diffbot's LLM-optimized mode."""
+
+    document_mapper: ExtractDocumentMapper | None = None
+    """Optional `(url, raw_response) -> Document` override."""
+
+    def _to_doc(self, url: str, raw: dict[str, Any]) -> Document:
+        mapper = self.document_mapper or _default_extract_mapper
+        return mapper(url, raw)
+
+    def lazy_load(self) -> Iterator[Document]:
+        """Yield one `Document` per URL, calling `extract` synchronously."""
+        with self._sync_db() as db:
+            for url in self.urls:
+                raw = db.extract(url, api=self.api, fmt=self.fmt)
+                yield self._to_doc(url, raw)
+
+    async def alazy_load(self) -> AsyncIterator[Document]:
+        """Yield one `Document` per URL, calling `extract` asynchronously."""
+        async with self._async_db() as db:
+            for url in self.urls:
+                raw = await db.extract(url, api=self.api, fmt=self.fmt)
+                yield self._to_doc(url, raw)
+
+
+class DiffbotCrawlLoader(_BaseDiffbotComponent, BaseLoader):
+    """Loader that drives a Diffbot crawl and yields a `Document` per URL.
+
+    By default `page_content` is the crawled URL itself (the crawl SDK only
+    yields URL events, not page contents). To fetch content per URL, chain
+    this with `DiffbotExtractLoader`.
+
+    Defaults `watch=True` so URL events are actually emitted — with
+    `watch=False` the SDK only yields a single JOB_CREATED event, which the
+    default mapper skips.
+    """
+
+    site: str
+    """Seed URL for the crawl."""
+
+    crawl_kwargs: dict[str, Any] = Field(default_factory=dict)
+    """Extra kwargs passed to `Diffbot.crawl(site, **crawl_kwargs)`.
+
+    Defaults to `{"watch": True}` if not overridden.
+    """
+
+    event_mapper: CrawlEventMapper | None = None
+    """Optional `(event) -> Document | None` override. Return None to skip."""
+
+    def _kwargs(self) -> dict[str, Any]:
+        kw = dict(self.crawl_kwargs)
+        kw.setdefault("watch", True)
+        return kw
+
+    def _map_event(self, event: CrawlEvent) -> Document | None:
+        mapper = self.event_mapper or _default_crawl_mapper
+        return mapper(event)
+
+    def lazy_load(self) -> Iterator[Document]:
+        """Drive a crawl synchronously and yield a `Document` per mapped event."""
+        with self._sync_db() as db:
+            for event in db.crawl(self.site, **self._kwargs()):
+                doc = self._map_event(event)
+                if doc is not None:
+                    yield doc
+
+    async def alazy_load(self) -> AsyncIterator[Document]:
+        """Drive a crawl asynchronously and yield a `Document` per mapped event."""
+        async with self._async_db() as db:
+            async for event in db.crawl(self.site, **self._kwargs()):
+                doc = self._map_event(event)
+                if doc is not None:
+                    yield doc