From 383613f9e8bfe4a8c066b07b3652008eb314d1aa Mon Sep 17 00:00:00 2001 From: Grady Dillon Date: Mon, 20 Jul 2026 16:46:39 -0400 Subject: [PATCH] Add WellMarkedRetriever and ruff lint - WellMarkedRetriever: the natural LangChain surface for the API's new /search endpoint. A query in, extracted result pages out as Documents, ready to drop into a RAG chain. Only successfully extracted results become Documents; a page that timed out or was blocked is skipped, since there is no content to embed. - Unit tests with the SDK client faked, matching the loader's test style (no network). - ruff added to the dev extras and to CI. Verified clean against the existing code before wiring it in, so it gates rather than breaks. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/ci.yml | 2 + langchain_wellmarked/__init__.py | 10 ++++- langchain_wellmarked/retrievers.py | 72 ++++++++++++++++++++++++++++++ pyproject.toml | 5 +++ tests/test_retrievers.py | 62 +++++++++++++++++++++++++ 5 files changed, 149 insertions(+), 2 deletions(-) create mode 100644 langchain_wellmarked/retrievers.py create mode 100644 tests/test_retrievers.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7759c8e..da83a56 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,5 +24,7 @@ jobs: cache: pip - run: python -m pip install --upgrade pip - run: pip install -e ".[dev]" + - name: Lint (ruff) + run: ruff check . - name: Test (pytest) run: pytest diff --git a/langchain_wellmarked/__init__.py b/langchain_wellmarked/__init__.py index 66028a5..02ea8d1 100644 --- a/langchain_wellmarked/__init__.py +++ b/langchain_wellmarked/__init__.py @@ -1,12 +1,18 @@ """Official LangChain integration for the WellMarked API. - from langchain_wellmarked import WellMarkedLoader + from langchain_wellmarked import WellMarkedLoader, WellMarkedRetriever + # Load specific URLs (or crawl a site) as Documents: loader = WellMarkedLoader("https://example.com/article") docs = loader.load() + # Or retrieve from the live web by query (search + extract): + retriever = WellMarkedRetriever(num_results=5) + docs = retriever.invoke("best open-source vector databases") + See https://wellmarked.io/docs for the full API reference. """ from langchain_wellmarked.document_loaders import WellMarkedLoader +from langchain_wellmarked.retrievers import WellMarkedRetriever -__all__ = ["WellMarkedLoader"] +__all__ = ["WellMarkedLoader", "WellMarkedRetriever"] diff --git a/langchain_wellmarked/retrievers.py b/langchain_wellmarked/retrievers.py new file mode 100644 index 0000000..e1ad9c7 --- /dev/null +++ b/langchain_wellmarked/retrievers.py @@ -0,0 +1,72 @@ +"""WellMarked retriever for LangChain.""" +from __future__ import annotations + +from typing import Any, List, Optional + +from langchain_core.callbacks import CallbackManagerForRetrieverRun +from langchain_core.documents import Document +from langchain_core.retrievers import BaseRetriever +from wellmarked import WellMarked + + +class WellMarkedRetriever(BaseRetriever): + """Retrieve web documents for a query via WellMarked search. + + Each retrieval searches the web, extracts the top results to clean Markdown, + and returns them as LangChain ``Document`` objects — search and extraction + in one round trip, no crawling or URL wrangling. This is the natural + LangChain surface for WellMarked: drop it into any RAG chain where the + context should come from the live web. + + Search requires a Pro plan or above (it is a paid feature); Free keys get + ``plan_not_supported``. + + Setup: + Install ``langchain-wellmarked`` and set ``WELLMARKED_API_KEY`` (or pass + ``api_key=``). Get a key at https://wellmarked.io. + + .. code-block:: bash + + pip install langchain-wellmarked + + Example: + .. code-block:: python + + from langchain_wellmarked import WellMarkedRetriever + + retriever = WellMarkedRetriever(num_results=5) + docs = retriever.invoke("best open-source vector databases") + docs[0].page_content # clean Markdown of a result page + docs[0].metadata # {"source": ..., "title": ..., "snippet": ...} + + Only successfully extracted results are returned; a result whose page timed + out or was blocked is skipped (its snippet stays in the API response but + there's no content to embed). + """ + + api_key: Optional[str] = None + """WellMarked API key (``wm_...``). Falls back to ``WELLMARKED_API_KEY``.""" + num_results: int = 5 + """How many results to fetch + extract. Clamped to 1..10 server-side.""" + render_js: bool = False + """Render JS-heavy result pages with a headless browser (Pro plan and above).""" + + def _get_relevant_documents( + self, query: str, *, run_manager: CallbackManagerForRetrieverRun, + ) -> List[Document]: + with WellMarked(api_key=self.api_key) as wm: + results = wm.search( + query, num_results=self.num_results, render_js=self.render_js, + ).results + + docs: List[Document] = [] + for r in results: + if not r.ok or not r.markdown: + continue + metadata: dict[str, Any] = {"source": r.url} + if r.title: + metadata["title"] = r.title + if r.snippet: + metadata["snippet"] = r.snippet + docs.append(Document(page_content=r.markdown, metadata=metadata)) + return docs diff --git a/pyproject.toml b/pyproject.toml index f9ab508..c3903cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,8 +33,13 @@ dependencies = [ [project.optional-dependencies] dev = [ "pytest>=7", + "ruff>=0.6", ] +[tool.ruff] +line-length = 100 +target-version = "py39" + [project.urls] Homepage = "https://wellmarked.io" Documentation = "https://wellmarked.io/docs" diff --git a/tests/test_retrievers.py b/tests/test_retrievers.py new file mode 100644 index 0000000..703b62e --- /dev/null +++ b/tests/test_retrievers.py @@ -0,0 +1,62 @@ +"""Unit tests for WellMarkedRetriever — the SDK client is faked, no network.""" +import pytest +from wellmarked import SearchResult, SearchResults + +import langchain_wellmarked.retrievers as mod +from langchain_wellmarked import WellMarkedRetriever + + +class FakeWellMarked: + """Stands in for wellmarked.WellMarked. Records calls, returns canned data.""" + + calls: list = [] + + def __init__(self, api_key=None, **kwargs): + FakeWellMarked.calls.append(("init", api_key)) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + pass + + def search(self, query, *, num_results=5, render_js=False): + FakeWellMarked.calls.append(("search", query, num_results, render_js)) + return SearchResults( + query=query, + results=[ + SearchResult( + url="https://a.test/1", status="ok", title="A", + snippet="s1", markdown="# A", + ), + SearchResult( + url="https://b.test/2", status="error", title="B", + snippet="s2", error="target_timeout", + ), + ], + request_id="req_1", + ) + + +@pytest.fixture(autouse=True) +def fake_client(monkeypatch): + FakeWellMarked.calls = [] + monkeypatch.setattr(mod, "WellMarked", FakeWellMarked) + + +def test_retriever_returns_ok_results_as_documents(): + docs = WellMarkedRetriever(api_key="wm_test", num_results=3).invoke("vector databases") + + assert ("init", "wm_test") in FakeWellMarked.calls + assert ("search", "vector databases", 3, False) in FakeWellMarked.calls + # The errored result is skipped — only extractable pages become Documents. + assert len(docs) == 1 + assert docs[0].page_content == "# A" + assert docs[0].metadata == { + "source": "https://a.test/1", "title": "A", "snippet": "s1", + } + + +def test_retriever_passes_render_js(): + WellMarkedRetriever(render_js=True).invoke("q") + assert ("search", "q", 5, True) in FakeWellMarked.calls