Add WellMarkedRetriever and ruff lint

.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

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"]

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

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"

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