From 96d8f8f862829d5814ca56fd30c370c8574f136d Mon Sep 17 00:00:00 2001 From: admin-raintree <277948009+admin-raintree@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:56:57 -0700 Subject: [PATCH 1/2] Add cited research answers from web search and local fetching `docpull research` and `docpull.research` search with a chosen backend, fetch results through the local pipeline, and rank paragraphs with the `docpull query` scorer. Answers are the best passages with exact spans and `[S#]` citations. An optional OpenAI-compatible model can rewrite them; the rewrite is kept only when every citation names a real source, and endpoints outside localhost require `--allow-remote-llm`. The compatible server adds Exa `/answer` and fills Tavily's `answer` when `include_answer` is set. Co-Authored-By: Claude Opus 5.5 --- docs/CHANGELOG.md | 6 + docs/cli-recipes.md | 29 +++- docs/surface-contract.md | 1 + src/docpull/cli.py | 4 + src/docpull/compat_api.py | 27 +++- src/docpull/research.py | 315 ++++++++++++++++++++++++++++++++++++++ src/docpull/server.py | 20 ++- src/docpull/surface.py | 2 +- tests/test_research.py | 163 ++++++++++++++++++++ 9 files changed, 561 insertions(+), 6 deletions(-) create mode 100644 src/docpull/research.py create mode 100644 tests/test_research.py diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 243f027..48baa07 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -19,6 +19,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `--include-content` fetches each result through the local pipeline. `docpull serve --compat --search-backend ...` also serves Tavily and Exa `/search` and Firecrawl `/v1/search` and `/v2/search` shapes. +- Answer questions from the web with `docpull research`. It searches, fetches + results locally, and returns ranked passages with exact spans and `[S#]` + citations. An optional OpenAI-compatible model, localhost unless + `--allow-remote-llm`, can rewrite them; the rewrite is kept only when every + citation names a real source. The compatible server adds Exa `/answer` and + Tavily `include_answer`. ### Removed diff --git a/docs/cli-recipes.md b/docs/cli-recipes.md index fa7c76f..6d3baf1 100644 --- a/docs/cli-recipes.md +++ b/docs/cli-recipes.md @@ -79,6 +79,30 @@ filter results by host. One include domain also narrows the upstream query. `docpull sync`, with HTTPS-only fetching, private-network blocking, and robots.txt handling, and attaches the Markdown. +## Research + +`docpull research` answers a question from the web with cited passages. It searches +with the chosen backend and fetches up to `--max-sources` results (default 5, +at most 10) through the local pipeline. It ranks paragraphs with the same scoring +as `docpull query` and returns the best `--max-passages` (default 8), at most three +per source. Each passage keeps its exact character span and source ID. + +```bash +docpull research "What is the GitHub REST API rate limit?" --backend searxng \ + --include-domain docs.github.com + +# Optional: rewrite the cited passages with a local OpenAI-compatible model +docpull research "What is the GitHub REST API rate limit?" --backend searxng \ + --llm-url http://localhost:11434/v1 --llm-model llama3.1 +``` + +Without a model, the answer is the ranked passages, each ending in a citation such +as `[S1]`. With `--llm-url` and `--llm-model`, the model may answer only from those +passages. Its answer is kept only when every citation names a real source. If it +cites nothing or names an unknown source, the cited passages are returned with a +warning. A model endpoint outside localhost needs `--allow-remote-llm`, because it +may cost money. `--llm-api-key-env NAME` sends a bearer key from the environment. + ## Local Extract API `docpull serve --compat` accepts the request shapes of common hosted extract @@ -98,7 +122,10 @@ docpull serve --compat --port 8765 Start the server with `--search-backend searxng` or `--search-backend marginalia` to also serve search shapes. Those are Tavily or Exa `POST /search`, chosen by the -request fields, and Firecrawl `POST /v1/search` and `/v2/search`. A client can +request fields, and Firecrawl `POST /v1/search` and `/v2/search`. The same server +answers Exa `POST /answer` and fills Tavily's `answer` when a search sets +`include_answer`. Both use `docpull research` ranking, plus the model from +`--llm-url` and `--llm-model` when one is configured. A client can pin a flavor by setting its base URL to `http://127.0.0.1:8765/tavily`, `/exa`, or `/firecrawl`. Without a search backend, the search routes return HTTP 501. diff --git a/docs/surface-contract.md b/docs/surface-contract.md index 1dd1faf..8a78b2a 100644 --- a/docs/surface-contract.md +++ b/docs/surface-contract.md @@ -43,6 +43,7 @@ DocPull targets capability alignment, not 1:1 flag parity. MCP should not mirror | List cached/indexed sources | Not a primary CLI command | Local filesystem/search helpers | `list_indexed` | MCP-focused | | Project lifecycle | `docpull init`, `add`, `install`, `deps`, `sources`, `sync`, `diff`, `status`, `history`, `review`, `release context-pack`, `watch` with explicit crawl bounds | `docpull.project` helpers | Not exposed yet | Core-aligned | | Search the web | `docpull search QUERY --backend searxng\|marginalia` with domain filters and `--include-content` | `search_web`, `SearchResponse`, `SearchHit` in `docpull.web_search` | Not exposed | CLI-first | +| Answer a question from the web with citations | `docpull research QUESTION --backend ...` with optional `--llm-url`/`--llm-model` | `research`, `answer_from_search`, `rank_evidence`, `LlmConfig` in `docpull.research` | Not exposed | CLI-first | | Search cached sources | SQLite/local search helpers where applicable | `search_sqlite_documents` | `grep_docs` | Adapted | | Query locked context with budgets, version filters, deduplication, and precise citations | `docpull query` | `query_context`, `ContextQueryError` | `context_query` | Core-aligned | | Read cached source by path/range | Filesystem responsibility | Filesystem responsibility | `read_doc` | MCP-specific | diff --git a/src/docpull/cli.py b/src/docpull/cli.py index 856f6e1..f40b0b6 100644 --- a/src/docpull/cli.py +++ b/src/docpull/cli.py @@ -1687,6 +1687,10 @@ def main(argv: list[str] | None = None) -> int: from .graph import run_graph_cli return run_graph_cli(raw_argv[1:]) + if raw_argv and raw_argv[0] == "research": + from .research import run_research_cli + + return run_research_cli(raw_argv[1:]) if raw_argv and raw_argv[0] == "search": from .web_search import run_search_cli diff --git a/src/docpull/compat_api.py b/src/docpull/compat_api.py index e6994bd..2b67f30 100644 --- a/src/docpull/compat_api.py +++ b/src/docpull/compat_api.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from typing import Any +from .research import LlmConfig, ResearchError, answer_from_search from .server import ASGIReceive, ASGIScope, ASGISend, _handle_lifespan, _json_response from .web_search import SearchResponse, WebSearchError @@ -73,10 +74,12 @@ def __init__( *, fetch: PageFetcher = fetch_page, search: SearchFn | None = None, + llm: LlmConfig | None = None, api_key: str | None = None, ) -> None: self.fetch = fetch self.search = search + self.llm = llm self.api_key = api_key async def __call__(self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend) -> None: @@ -118,7 +121,7 @@ async def __call__(self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend) except SearchUnavailableError as err: await _json_response(send, 501, {"error": str(err)}) return - except WebSearchError as err: + except (WebSearchError, ResearchError) as err: await _json_response(send, 502, {"error": str(err)}) return await _json_response(send, 200, payload) @@ -222,13 +225,17 @@ class SearchUnavailableError(RuntimeError): async def _tavily_search(app: CompatASGIApp, body: dict[str, Any]) -> dict[str, Any]: started = time.perf_counter() raw = body.get("include_raw_content", False) + wants_answer = bool(body.get("include_answer", False)) response = await app.run_search( body.get("query"), body.get("max_results", 5), body.get("include_domains"), body.get("exclude_domains"), - bool(raw), + bool(raw) or wants_answer, ) + answer = None + if wants_answer: + answer = (await answer_from_search(response.query, response, llm=app.llm)).answer or None results = [ { "title": hit.title, @@ -241,7 +248,7 @@ async def _tavily_search(app: CompatASGIApp, body: dict[str, Any]) -> dict[str, ] return { "query": response.query, - "answer": None, + "answer": answer, "images": [], "follow_up_questions": None, "results": results, @@ -275,6 +282,19 @@ async def _exa_search(app: CompatASGIApp, body: dict[str, Any]) -> dict[str, Any return {"requestId": uuid.uuid4().hex, "resolvedSearchType": "keyword", "results": results} +async def _exa_answer(app: CompatASGIApp, body: dict[str, Any]) -> dict[str, Any]: + response = await app.run_search(body.get("query"), 5, None, None, True) + result = await answer_from_search(response.query, response, llm=app.llm) + include_text = bool(body.get("text", False)) + citations = [] + for source in result.sources: + citation: dict[str, Any] = {"id": source["url"], "url": source["url"], "title": source["title"]} + if include_text: + citation["text"] = "\n\n".join(e.text for e in result.evidence if e.source_id == source["id"]) + citations.append(citation) + return {"requestId": uuid.uuid4().hex, "answer": result.answer, "citations": citations} + + async def _firecrawl_search(app: CompatASGIApp, body: dict[str, Any], *, version: int) -> dict[str, Any]: options = body.get("scrapeOptions") formats = options.get("formats") if isinstance(options, dict) else None @@ -306,6 +326,7 @@ async def _firecrawl_search_v2(app: CompatASGIApp, body: dict[str, Any]) -> dict "/v1/scrape": _firecrawl_scrape, "/v2/scrape": _firecrawl_scrape, "/search": _tavily_search, + "/answer": _exa_answer, "/v1/search": _firecrawl_search_v1, "/v2/search": _firecrawl_search_v2, } diff --git a/src/docpull/research.py b/src/docpull/research.py new file mode 100644 index 0000000..a762483 --- /dev/null +++ b/src/docpull/research.py @@ -0,0 +1,315 @@ +"""Cited answers built from web search, local fetching, and passage ranking. + +``research`` searches with a backend the caller names, fetches the top results +through the local pipeline, and ranks paragraphs with the same scoring as +``docpull query``. By default the answer is the best passages, each cited. An +OpenAI-compatible model can rewrite them into prose, but only if every citation +it returns points at a real source; otherwise the cited passages are returned. +""" + +from __future__ import annotations + +import argparse +import asyncio +import ipaddress +import json +import os +import re +import sys +from dataclasses import asdict, dataclass, field +from typing import Any +from urllib.parse import urlsplit + +import aiohttp + +from .context_query import _STOPWORDS, _spans +from .pack_tools import _search_score, _term_count +from .web_search import BACKENDS, SearchHit, SearchResponse, WebSearchError, search_web + +MAX_SOURCES = 10 +MAX_PASSAGES_PER_SOURCE = 3 +PASSAGE_CHARS = 600 +_CITATION_RE = re.compile(r"\[S(\d+)\]") +_LLM_TIMEOUT = aiohttp.ClientTimeout(total=90) + + +class ResearchError(RuntimeError): + """Research could not run with the given configuration.""" + + +@dataclass(frozen=True) +class LlmConfig: + """An OpenAI-compatible chat completions endpoint, such as Ollama or llama.cpp.""" + + base_url: str + model: str + api_key: str | None = None + allow_remote: bool = False + + +@dataclass(frozen=True) +class Evidence: + source_id: str + url: str + title: str + text: str + char_start: int + char_end: int + score: int + + +@dataclass(frozen=True) +class ResearchResult: + query: str + answer: str + mode: str + sources: list[dict[str, str]] + evidence: list[Evidence] + search: SearchResponse + warnings: list[str] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + payload = asdict(self) + payload["search"] = {"backend": self.search.backend, "license": self.search.license} + return payload + + +async def research( + query: str, + *, + backend: str, + max_sources: int = 5, + max_passages: int = 8, + include_domains: list[str] | None = None, + exclude_domains: list[str] | None = None, + searxng_url: str | None = None, + marginalia_key: str | None = None, + llm: LlmConfig | None = None, +) -> ResearchResult: + if not 1 <= max_sources <= MAX_SOURCES: + raise ResearchError(f"max_sources must be between 1 and {MAX_SOURCES}") + if llm is not None: + _check_llm(llm) + try: + search = await search_web( + query, + backend=backend, + max_results=max_sources, + include_domains=include_domains, + exclude_domains=exclude_domains, + include_content=True, + searxng_url=searxng_url, + marginalia_key=marginalia_key, + ) + except WebSearchError as err: + raise ResearchError(str(err)) from err + return await answer_from_search(query, search, max_passages=max_passages, llm=llm) + + +async def answer_from_search( + query: str, + search: SearchResponse, + *, + max_passages: int = 8, + llm: LlmConfig | None = None, +) -> ResearchResult: + """Rank fetched search results into cited evidence and an answer.""" + warnings = list(search.warnings) + failed = [hit.url for hit in search.hits if hit.content_error] + if failed: + warnings.append(f"{len(failed)} result(s) could not be fetched and were skipped") + evidence = rank_evidence(query, search.hits, max_passages=max_passages) + first = {item.source_id: item for item in reversed(evidence)} + sources = [ + {"id": source_id, "url": first[source_id].url, "title": first[source_id].title} + for source_id in dict.fromkeys(item.source_id for item in evidence) + ] + if not evidence: + warnings.append("No fetched passage matched the question") + return ResearchResult(search.query, "", "extractive", [], [], search, warnings) + + extractive = _extractive_answer(evidence) + if llm is None: + return ResearchResult(search.query, extractive, "extractive", sources, evidence, search, warnings) + try: + answer = await _synthesize(query, evidence, llm) + except ResearchError as err: + warnings.append(f"Model answer discarded: {err}") + return ResearchResult(search.query, extractive, "extractive", sources, evidence, search, warnings) + return ResearchResult(search.query, answer, "synthesized", sources, evidence, search, warnings) + + +def rank_evidence(query: str, hits: list[SearchHit], *, max_passages: int = 8) -> list[Evidence]: + """Score paragraphs across fetched hits; source IDs follow search rank.""" + terms = sorted(set(re.findall(r"[^\W_]+(?:[_-][^\W_]+)*", query.casefold())) - _STOPWORDS) + if not terms: + raise ResearchError("The question has no searchable terms") + candidates: list[tuple[int, int, Evidence]] = [] + for rank, hit in enumerate(hits, start=1): + content = hit.content or "" + scored: list[Evidence] = [] + for start, end in _spans(content): + raw = content[start:end] + start += len(raw) - len(raw.lstrip()) + end -= len(raw) - len(raw.rstrip()) + text = content[start:end] + if not text or not any(_term_count(text, term) for term in terms): + continue + score, _ = _search_score( + query_terms=terms, + phrase=query.strip().lower(), + title=hit.title, + heading="", + url=hit.url, + content=text, + ) + scored.append(Evidence(f"S{rank}", hit.url, hit.title, text, start, end, score)) + scored.sort(key=lambda item: (-item.score, item.char_start)) + candidates.extend((item.score, rank, item) for item in scored[:MAX_PASSAGES_PER_SOURCE]) + candidates.sort(key=lambda entry: (-entry[0], entry[1], entry[2].char_start)) + chosen = [item for _, _, item in candidates[:max_passages]] + # Renumber sources densely in order of first use so citations read S1, S2, ... + renumber = { + old: f"S{index}" for index, old in enumerate(dict.fromkeys(e.source_id for e in chosen), start=1) + } + return [ + Evidence(renumber[e.source_id], e.url, e.title, e.text, e.char_start, e.char_end, e.score) + for e in chosen + ] + + +def _extractive_answer(evidence: list[Evidence]) -> str: + lines = [] + for item in evidence: + text = " ".join(item.text.split()) + if len(text) > PASSAGE_CHARS: + text = text[: PASSAGE_CHARS - 1].rsplit(" ", 1)[0] + "…" + lines.append(f"- {text} [{item.source_id}]") + return "\n".join(lines) + + +def _check_llm(llm: LlmConfig) -> None: + parts = urlsplit(llm.base_url) + if parts.scheme not in {"http", "https"} or not parts.hostname: + raise ResearchError("llm base URL must be an http(s) URL") + if not llm.model.strip(): + raise ResearchError("llm model must not be empty") + if not llm.allow_remote and not _is_loopback(parts.hostname): + raise ResearchError("A non-localhost model endpoint requires --allow-remote-llm; it may cost money") + + +def _is_loopback(host: str) -> bool: + if host.lower() == "localhost": + return True + try: + return ipaddress.ip_address(host.strip("[]")).is_loopback + except ValueError: + return False + + +async def _synthesize(query: str, evidence: list[Evidence], llm: LlmConfig) -> str: + numbered = "\n\n".join(f"[{item.source_id}] {item.title} ({item.url})\n{item.text}" for item in evidence) + messages = [ + { + "role": "system", + "content": ( + "Answer only from the numbered evidence. Cite every claim with its source tag, " + "for example [S1]. If the evidence does not answer the question, say so." + ), + }, + {"role": "user", "content": f"Question: {query}\n\nEvidence:\n{numbered}"}, + ] + headers = {"Authorization": f"Bearer {llm.api_key}"} if llm.api_key else {} + url = llm.base_url.rstrip("/") + "/chat/completions" + body = {"model": llm.model, "messages": messages, "temperature": 0, "max_tokens": 800} + try: + async with ( + aiohttp.ClientSession(timeout=_LLM_TIMEOUT) as session, + session.post(url, json=body, headers=headers) as response, + ): + if response.status >= 400: + raise ResearchError(f"model endpoint returned HTTP {response.status}") + payload = await response.json(content_type=None) + except (aiohttp.ClientError, TimeoutError, json.JSONDecodeError) as err: + raise ResearchError(f"model request failed: {type(err).__name__}") from err + try: + answer = str(payload["choices"][0]["message"]["content"]).strip() + except (KeyError, IndexError, TypeError) as err: + raise ResearchError("model returned no answer") from err + valid = {item.source_id for item in evidence} + cited = {f"S{number}" for number in _CITATION_RE.findall(answer)} + if not cited: + raise ResearchError("answer cites no sources") + if not cited <= valid: + raise ResearchError(f"answer cites unknown sources: {', '.join(sorted(cited - valid))}") + return answer + + +def create_research_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="docpull research", + description="Answer a question with cited passages from web search and local fetching", + ) + parser.add_argument("question") + parser.add_argument("--backend", choices=BACKENDS, required=True) + parser.add_argument("--max-sources", type=int, default=5) + parser.add_argument("--max-passages", type=int, default=8) + parser.add_argument("--include-domain", action="append", default=[], dest="include_domains") + parser.add_argument("--exclude-domain", action="append", default=[], dest="exclude_domains") + parser.add_argument("--searxng-url") + parser.add_argument("--llm-url", help="OpenAI-compatible base URL, e.g. http://localhost:11434/v1") + parser.add_argument("--llm-model", help="Model name for --llm-url") + parser.add_argument( + "--llm-api-key-env", metavar="NAME", help="Environment variable holding the model key" + ) + parser.add_argument( + "--allow-remote-llm", + action="store_true", + help="Allow a non-localhost model endpoint, which may cost money", + ) + parser.add_argument("--json", action="store_true", dest="json_output") + return parser + + +def llm_config_from_args(args: argparse.Namespace) -> LlmConfig | None: + if not args.llm_url and not args.llm_model: + return None + if not (args.llm_url and args.llm_model): + raise ResearchError("--llm-url and --llm-model must be given together") + api_key = os.environ.get(args.llm_api_key_env) if args.llm_api_key_env else None + if args.llm_api_key_env and not api_key: + raise ResearchError(f"Environment variable {args.llm_api_key_env} is empty or unset") + return LlmConfig(args.llm_url, args.llm_model, api_key=api_key, allow_remote=args.allow_remote_llm) + + +def run_research_cli(argv: list[str] | None = None) -> int: + args = create_research_parser().parse_args(argv) + try: + result = asyncio.run( + research( + args.question, + backend=args.backend, + max_sources=args.max_sources, + max_passages=args.max_passages, + include_domains=args.include_domains, + exclude_domains=args.exclude_domains, + searxng_url=args.searxng_url, + llm=llm_config_from_args(args), + ) + ) + except ResearchError as err: + print(f"docpull research: {err}", file=sys.stderr) + return 1 + if args.json_output: + print(json.dumps(result.to_dict(), ensure_ascii=False, indent=2)) + return 0 + for warning in result.warnings: + print(f"warning: {warning}", file=sys.stderr) + print(result.answer or "No answer: no fetched passage matched the question.") + if result.sources: + print("\nSources:") + for source in result.sources: + print(f"[{source['id']}] {source['title'] or source['url']} — {source['url']}") + if result.search.license: + print(f"\nSearch results license: {result.search.license}") + return 0 diff --git a/src/docpull/server.py b/src/docpull/server.py index c987c6d..d4690bd 100644 --- a/src/docpull/server.py +++ b/src/docpull/server.py @@ -131,6 +131,16 @@ def create_serve_parser() -> argparse.ArgumentParser: help="With --compat, serve /search routes through this backend", ) parser.add_argument("--searxng-url", help="With --search-backend searxng, the SearXNG base URL") + parser.add_argument("--llm-url", help="With --compat, an OpenAI-compatible base URL for answers") + parser.add_argument("--llm-model", help="Model name for --llm-url") + parser.add_argument( + "--llm-api-key-env", metavar="NAME", help="Environment variable holding the model key" + ) + parser.add_argument( + "--allow-remote-llm", + action="store_true", + help="Allow a non-localhost model endpoint, which may cost money", + ) parser.add_argument( "--api-key-env", metavar="NAME", @@ -173,6 +183,7 @@ def run_serve_cli(argv: list[str] | None = None) -> int: def _run_compat(args: argparse.Namespace, console: Console) -> int: from .compat_api import CompatASGIApp + from .research import ResearchError, _check_llm, llm_config_from_args from .web_search import search_web api_key = None @@ -184,6 +195,12 @@ def _run_compat(args: argparse.Namespace, console: Console) -> int: raise PackServerError(f"Environment variable {args.api_key_env} is empty or unset.") if api_key is None and not _is_loopback_host(args.host): raise PackServerError("A non-localhost --compat server requires --api-key-env.") + try: + llm = llm_config_from_args(args) + if llm is not None: + _check_llm(llm) + except ResearchError as err: + raise PackServerError(str(err)) from err import uvicorn except PackServerError as err: console.print("[red]Serve error:[/red] " + escape(str(err))) @@ -208,7 +225,8 @@ async def search( searxng_url=args.searxng_url, ) - uvicorn.run(CompatASGIApp(search=search, api_key=api_key), host=args.host, port=args.port) + app = CompatASGIApp(search=search, llm=llm, api_key=api_key) + uvicorn.run(app, host=args.host, port=args.port) return 0 diff --git a/src/docpull/surface.py b/src/docpull/surface.py index ebe7651..b6f74e2 100644 --- a/src/docpull/surface.py +++ b/src/docpull/surface.py @@ -36,6 +36,7 @@ class CliCommand: CliCommand("parse", "Parse local PDF/DOCX/text files into a v3 pack"), CliCommand("render", "Render one public URL or check optional browser backends"), CliCommand("search", "Search the web through self-hosted SearXNG or Marginalia"), + CliCommand("research", "Answer a question with cited passages from web search and local fetching"), CliCommand("openapi-pack", "Build a local v3 pack from an OpenAPI spec"), CliCommand("feed-pack", "Build a local v3 pack from an RSS, Atom, JSON Feed, or feed page"), CliCommand("paper-pack", "Build a local v3 pack from papers and scholarly metadata"), @@ -77,7 +78,6 @@ class CliCommand: "batch", "crawl", "extract", - "research", "answer", "entities", "brief", diff --git a/tests/test_research.py b/tests/test_research.py new file mode 100644 index 0000000..691ca71 --- /dev/null +++ b/tests/test_research.py @@ -0,0 +1,163 @@ +"""Tests for cited research answers and the compatible answer routes.""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import pytest +from aiohttp import web + +from docpull import research as research_module +from docpull.compat_api import CompatASGIApp +from docpull.research import ( + LlmConfig, + ResearchError, + _check_llm, + answer_from_search, + rank_evidence, + run_research_cli, +) +from docpull.web_search import SearchHit, SearchResponse +from tests.test_compat_api import _post + +ALPHA = "# Rate limits\n\nThe API allows 100 requests per minute per key.\n\nUnrelated footer text.\n" +BETA = "Intro.\n\nExceeding the rate limit returns HTTP 429 with a Retry-After header.\n" + + +def _search(*, beta_error: str | None = None) -> SearchResponse: + return SearchResponse( + query="rate limit per minute", + backend="fake", + hits=[ + SearchHit(url="https://docs.example.com/limits", title="Limits", content=ALPHA), + SearchHit(url="https://empty.example.com", title="Empty", content="Nothing relevant here."), + SearchHit( + url="https://blog.example.com/429", + title="429", + content=None if beta_error else BETA, + content_error=beta_error, + ), + ], + ) + + +def test_rank_evidence_keeps_exact_spans_and_dense_source_ids() -> None: + evidence = rank_evidence("rate limit per minute", _search().hits) + + assert {item.source_id for item in evidence} == {"S1", "S2"} + assert all(item.url != "https://empty.example.com" for item in evidence) + for item in evidence: + source = ALPHA if item.url.endswith("/limits") else BETA + assert source[item.char_start : item.char_end] == item.text + with pytest.raises(ResearchError, match="no searchable terms"): + rank_evidence("the a of", _search().hits) + + +@pytest.mark.asyncio +async def test_extractive_answer_cites_every_passage() -> None: + result = await answer_from_search("rate limit per minute", _search()) + + assert result.mode == "extractive" + assert all(line.rstrip().endswith(("[S1]", "[S2]")) for line in result.answer.splitlines()) + assert [source["id"] for source in result.sources] == ["S1", "S2"] + + partial = await answer_from_search("rate limit", _search(beta_error="HTTP 404")) + assert any("could not be fetched" in warning for warning in partial.warnings) + + empty = await answer_from_search("zyzzyva", _search()) + assert empty.answer == "" and empty.sources == [] + assert "No fetched passage matched the question" in empty.warnings + + +@pytest.fixture +async def fake_model(): + replies: list[str] = [] + requests: list[dict[str, Any]] = [] + + async def completions(request: web.Request) -> web.Response: + requests.append(await request.json()) + return web.json_response({"choices": [{"message": {"content": replies.pop(0)}}]}) + + app = web.Application() + app.router.add_post("/v1/chat/completions", completions) + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + port = site._server.sockets[0].getsockname()[1] # type: ignore[union-attr] + yield LlmConfig(f"http://127.0.0.1:{port}/v1", "local-model"), replies, requests + await runner.cleanup() + + +@pytest.mark.asyncio +async def test_model_answer_is_kept_only_with_valid_citations(fake_model: Any) -> None: + llm, replies, requests = fake_model + replies.extend( + [ + "The limit is 100 requests per minute [S1]; excess returns 429 [S2].", + "The limit is 100 per minute [S9].", + "The limit is 100 per minute.", + ] + ) + + good = await answer_from_search("rate limit per minute", _search(), llm=llm) + assert good.mode == "synthesized" and "[S2]" in good.answer + assert requests[0]["model"] == "local-model" and requests[0]["temperature"] == 0 + assert "[S1] Limits" in requests[0]["messages"][1]["content"] + + unknown = await answer_from_search("rate limit per minute", _search(), llm=llm) + assert unknown.mode == "extractive" + assert any("unknown sources: S9" in warning for warning in unknown.warnings) + + uncited = await answer_from_search("rate limit per minute", _search(), llm=llm) + assert uncited.mode == "extractive" + assert any("cites no sources" in warning for warning in uncited.warnings) + + +def test_remote_model_endpoints_require_opt_in() -> None: + _check_llm(LlmConfig("http://localhost:11434/v1", "m")) + with pytest.raises(ResearchError, match="allow-remote-llm"): + _check_llm(LlmConfig("https://api.example.com/v1", "m")) + _check_llm(LlmConfig("https://api.example.com/v1", "m", allow_remote=True)) + with pytest.raises(ResearchError, match="http"): + _check_llm(LlmConfig("file:///etc/passwd", "m")) + + +@pytest.mark.asyncio +async def test_research_cli_json(monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]) -> None: + async def fake_search_web(query: str, **kwargs: Any) -> SearchResponse: + assert kwargs["include_content"] is True + return _search() + + monkeypatch.setattr(research_module, "search_web", fake_search_web) + argv = ["rate limit per minute", "--backend", "marginalia", "--json"] + assert await asyncio.to_thread(run_research_cli, argv) == 0 + payload = json.loads(capsys.readouterr().out) + assert payload["mode"] == "extractive" + assert payload["sources"][0]["id"] == "S1" + assert payload["search"] == {"backend": "fake", "license": None} + + assert await asyncio.to_thread(run_research_cli, [*argv[:-1], "--llm-url", "http://localhost:1/v1"]) == 1 + + +async def _search_fn(query: str, max_results: int, include: list[str], exclude: list[str], content: bool): + assert content is True + return _search() + + +@pytest.mark.asyncio +async def test_exa_answer_and_tavily_include_answer_shapes() -> None: + app = CompatASGIApp(search=_search_fn) + + status, exa = await _post(app, "/answer", {"query": "rate limit per minute", "text": True}) + assert status == 200 + assert "[S1]" in exa["answer"] + assert exa["citations"][0]["url"] == "https://docs.example.com/limits" + assert "100 requests per minute" in exa["citations"][0]["text"] + + status, tavily = await _post(app, "/search", {"query": "rate limit per minute", "include_answer": True}) + assert status == 200 + assert "[S1]" in tavily["answer"] + assert tavily["results"][0]["raw_content"] is None From 49498eb02313e544e4984437f4dbf9379642506e Mon Sep 17 00:00:00 2001 From: admin-raintree <277948009+admin-raintree@users.noreply.github.com> Date: Mon, 12 Oct 2026 13:09:03 -0700 Subject: [PATCH 2/2] Run async crawl, task-run, and research jobs on the compatible server Add Firecrawl `/v1/crawl` and `/v2/crawl` (create, poll, cancel), Parallel `/v1/tasks/runs` with `/result`, and Tavily `/research` to `docpull serve --compat`. Crawls use local discovery and fetching with Firecrawl's regex path filters; task runs and research tasks run `docpull research` and return cited passages. Jobs run as background tasks, live in memory, expire after 24 hours, and are capped at 200. A done-callback marks jobs finished even when they are cancelled before starting. Co-Authored-By: Claude Opus 5.5 --- docs/CHANGELOG.md | 5 + docs/cli-recipes.md | 11 ++ src/docpull/compat_api.py | 104 ++++++++++- src/docpull/compat_jobs.py | 355 +++++++++++++++++++++++++++++++++++++ tests/test_compat_jobs.py | 176 ++++++++++++++++++ 5 files changed, 647 insertions(+), 4 deletions(-) create mode 100644 src/docpull/compat_jobs.py create mode 100644 tests/test_compat_jobs.py diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 48baa07..8323f19 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -25,6 +25,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `--allow-remote-llm`, can rewrite them; the rewrite is kept only when every citation names a real source. The compatible server adds Exa `/answer` and Tavily `include_answer`. +- Run async jobs on the compatible server. Firecrawl `/v1/crawl` and `/v2/crawl` + crawl locally with regex path filters and cancellation. Parallel + `/v1/tasks/runs` (with `/result`) and Tavily `/research` run + `docpull research`. Jobs are kept in memory, expire after 24 hours, and are + capped at 200. ### Removed diff --git a/docs/cli-recipes.md b/docs/cli-recipes.md index 6d3baf1..f7d0dde 100644 --- a/docs/cli-recipes.md +++ b/docs/cli-recipes.md @@ -130,6 +130,17 @@ pin a flavor by setting its base URL to `http://127.0.0.1:8765/tavily`, `/exa`, or `/firecrawl`. Without a search backend, the search routes return HTTP 501. +The server also runs async jobs in the shapes those APIs use: + +| Create | Poll | Notes | +| --- | --- | --- | +| Firecrawl `POST /v1/crawl` or `/v2/crawl` with `url`, `limit` (default 50, at most 500), `maxDiscoveryDepth`, `includePaths`, `excludePaths` | `GET /v{1,2}/crawl/{id}`; `DELETE` cancels | Discovery and fetching use the local crawler. Path filters are regular expressions on the URL path, as in Firecrawl. `creditsUsed` is always 0 | +| Parallel `POST /v1/tasks/runs` with `input` | `GET /v1/tasks/runs/{run_id}`; `GET .../result?timeout=` waits up to 600 seconds | Runs `docpull research`. `output.basis[0].citations` lists each source with its passages as `excerpts`. Needs `--search-backend` | +| Tavily `POST /research` with `input` | `GET /research/{request_id}` returns 202 until done | Runs `docpull research`. Returns `content` and `sources`. Needs `--search-backend` | + +Jobs live in memory, so they are lost when the server stops. Finished jobs expire +after 24 hours, and the server keeps at most 200 jobs. + Every page goes through the same fetcher as `docpull sync`. That means HTTPS-only fetching, private-network blocking, and robots.txt handling. Nothing is sent to the hosted services. A request takes at most 20 URLs. The server binds to diff --git a/src/docpull/compat_api.py b/src/docpull/compat_api.py index 2b67f30..d176695 100644 --- a/src/docpull/compat_api.py +++ b/src/docpull/compat_api.py @@ -11,12 +11,28 @@ import asyncio import hmac import json +import re import time import uuid from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import Any +from .compat_jobs import ( + RESULT_WAIT_SECONDS, + Crawler, + JobRequestError, + JobStore, + cancel_crawl, + crawl_site, + crawl_status, + create_crawl, + create_research, + create_task_run, + research_status, + task_run, + task_run_result, +) from .research import LlmConfig, ResearchError, answer_from_search from .server import ASGIReceive, ASGIScope, ASGISend, _handle_lifespan, _json_response from .web_search import SearchResponse, WebSearchError @@ -49,7 +65,6 @@ class CompatRequestError(ValueError): async def fetch_page(url: str) -> Page: """Fetch one URL to Markdown in memory with the default local pipeline.""" - from .conversion.special_cases import _split_markdown_frontmatter from .core.fetcher import Fetcher from .models.config import DocpullConfig @@ -58,6 +73,13 @@ async def fetch_page(url: str) -> Page: ctx = await fetcher.fetch_one(url, save=False) except Exception as err: # noqa: BLE001 - surfaced per URL, like the hosted APIs return Page(url=url, error=str(err) or type(err).__name__) + return page_from_context(url, ctx) + + +def page_from_context(url: str, ctx: Any) -> Page: + """Convert a fetched ``PageContext`` into a served page body.""" + from .conversion.special_cases import _split_markdown_frontmatter + if ctx.error or ctx.should_skip or not (ctx.markdown or "").strip(): reason = ctx.error or ctx.skip_reason or "No extractable content" return Page(url=url, status_code=ctx.status_code, error=reason) @@ -76,7 +98,10 @@ def __init__( search: SearchFn | None = None, llm: LlmConfig | None = None, api_key: str | None = None, + crawl: Crawler = crawl_site, ) -> None: + self.jobs = JobStore() + self.crawl = crawl self.fetch = fetch self.search = search self.llm = llm @@ -97,11 +122,25 @@ async def __call__(self, scope: ASGIScope, receive: ASGIReceive, send: ASGISend) if path == "/health" and scope.get("method") == "GET": await _json_response(send, 200, {"status": "ok", "endpoints": sorted(_ROUTES)}) return + method = str(scope.get("method") or "GET") + if _is_job_route(method, path): + if not self._authorized(scope): + await _json_response(send, 401, {"error": "Missing or invalid API key"}) + return + try: + body = await _read_body(receive) if method == "POST" else {} + status, payload = await _dispatch_job(self, method, path, body, scope, flavor) + except (CompatRequestError, JobRequestError) as err: + status, payload = 400, {"error": str(err)} + except SearchUnavailableError as err: + status, payload = 501, {"error": str(err)} + await _json_response(send, status, payload) + return handler = _ROUTES.get(path) if handler is None: await _json_response(send, 404, {"error": "Not found", "path": path}) return - if scope.get("method") != "POST": + if method != "POST": await _json_response(send, 405, {"error": "Use POST"}) return if not self._authorized(scope): @@ -137,11 +176,15 @@ def _authorized(self, scope: ASGIScope) -> bool: ) return hmac.compare_digest(supplied.encode(), self.api_key.encode()) + def require_search(self) -> None: + if self.search is None: + raise SearchUnavailableError("Search is not configured; start the server with --search-backend") + async def run_search( self, query: Any, max_results: Any, include: Any, exclude: Any, content: bool ) -> SearchResponse: - if self.search is None: - raise SearchUnavailableError("Search is not configured; start the server with --search-backend") + self.require_search() + assert self.search is not None if not isinstance(query, str) or not query.strip(): raise CompatRequestError("query is required") if isinstance(max_results, bool) or not isinstance(max_results, int) or not 1 <= max_results <= 20: @@ -332,6 +375,59 @@ async def _firecrawl_search_v2(app: CompatASGIApp, body: dict[str, Any]) -> dict } +_CRAWL_ID = re.compile(r"^/v[12]/crawl/([\w-]+)$") +_RUN_ID = re.compile(r"^/v1/tasks/runs/([\w-]+)(/result)?$") +_RESEARCH_ID = re.compile(r"^/research/([\w-]+)$") +_JOB_CREATE_PATHS = frozenset({"/v1/crawl", "/v2/crawl", "/v1/tasks/runs", "/research"}) + + +def _is_job_route(method: str, path: str) -> bool: + if method == "POST": + return path in _JOB_CREATE_PATHS + if method == "DELETE": + return bool(_CRAWL_ID.match(path)) + return method == "GET" and any(pattern.match(path) for pattern in (_CRAWL_ID, _RUN_ID, _RESEARCH_ID)) + + +async def _dispatch_job( + app: CompatASGIApp, method: str, path: str, body: dict[str, Any], scope: ASGIScope, flavor: str | None +) -> tuple[int, dict[str, Any]]: + if method == "POST": + if path == "/v1/tasks/runs": + return await create_task_run(app, body) + if path == "/research": + return await create_research(app, body) + return await create_crawl(app, body, _origin(scope, flavor) + path) + if match := _CRAWL_ID.match(path): + return cancel_crawl(app, match.group(1)) if method == "DELETE" else crawl_status(app, match.group(1)) + if match := _RUN_ID.match(path): + if match.group(2): + return await task_run_result(app, match.group(1), _query_timeout(scope)) + return task_run(app, match.group(1)) + match = _RESEARCH_ID.match(path) + assert match is not None + return research_status(app, match.group(1)) + + +def _origin(scope: ASGIScope, flavor: str | None) -> str: + headers = { + key.decode("latin-1").lower(): value.decode("latin-1") for key, value in scope.get("headers", []) + } + host = headers.get("host") or "127.0.0.1" + prefix = f"/{flavor}" if flavor else "" + return f"{scope.get('scheme', 'http')}://{host}{prefix}" + + +def _query_timeout(scope: ASGIScope) -> float: + from urllib.parse import parse_qs + + values = parse_qs(bytes(scope.get("query_string") or b"").decode()).get("timeout") + try: + return max(0.0, min(float(values[0]), RESULT_WAIT_SECONDS)) if values else RESULT_WAIT_SECONDS + except ValueError: + return RESULT_WAIT_SECONDS + + def _domains(value: Any) -> list[str]: if value is None: return [] diff --git a/src/docpull/compat_jobs.py b/src/docpull/compat_jobs.py new file mode 100644 index 0000000..891e68b --- /dev/null +++ b/src/docpull/compat_jobs.py @@ -0,0 +1,355 @@ +"""Async job routes for the compatible server: crawl, task runs, and research. + +Jobs run as background tasks in the server process and live in memory, so they +are lost when the server stops. +""" + +from __future__ import annotations + +import asyncio +import re +import time +import uuid +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Any + +from .research import ResearchResult, answer_from_search + +if TYPE_CHECKING: + from .compat_api import CompatASGIApp, Page + +# ponytail: in-memory store, capped and expiring; persist jobs if restarts must keep them. +MAX_JOBS = 200 +JOB_TTL = timedelta(hours=24) +DEFAULT_CRAWL_PAGES = 50 +MAX_CRAWL_PAGES = 500 +CRAWL_CONCURRENCY = 4 +RESULT_WAIT_SECONDS = 600 + +Crawler = Callable[[str, int, int, list[str], list[str], Callable[["Page"], None]], Awaitable[int]] + + +class JobRequestError(ValueError): + """Invalid job request.""" + + +@dataclass +class Job: + id: str + kind: str + request: dict[str, Any] + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + modified_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + status: str = "queued" + pages: list[Page] = field(default_factory=list) + total: int = 0 + result: ResearchResult | None = None + error: str | None = None + started: float = field(default_factory=time.perf_counter) + finished: float | None = None + done: asyncio.Event = field(default_factory=asyncio.Event) + task: asyncio.Task[None] | None = None + + def touch(self, status: str) -> None: + self.status = status + self.modified_at = datetime.now(timezone.utc) + + @property + def elapsed(self) -> float: + return round((self.finished or time.perf_counter()) - self.started, 3) + + +class JobStore: + def __init__(self) -> None: + self._jobs: dict[str, Job] = {} + + def add(self, kind: str, request: dict[str, Any], prefix: str) -> Job: + self._expire() + if len(self._jobs) >= MAX_JOBS: + raise JobRequestError(f"Too many running jobs; at most {MAX_JOBS} are kept") + job = Job(id=f"{prefix}{uuid.uuid4().hex}", kind=kind, request=request) + self._jobs[job.id] = job + return job + + def get(self, job_id: str, kind: str) -> Job | None: + self._expire() + job = self._jobs.get(job_id) + return job if job and job.kind == kind else None + + def _expire(self) -> None: + now = datetime.now(timezone.utc) + for job_id, job in list(self._jobs.items()): + if job.done.is_set() and now - job.modified_at > JOB_TTL: + del self._jobs[job_id] + finished = [job for job in self._jobs.values() if job.done.is_set()] + for job in sorted(finished, key=lambda item: item.modified_at)[ + : max(0, len(self._jobs) - MAX_JOBS + 1) + ]: + del self._jobs[job.id] + + +def start(job: Job, work: Callable[[Job], Awaitable[None]]) -> None: + async def run() -> None: + job.touch("running") + try: + await work(job) + job.touch("completed") + except Exception as err: # noqa: BLE001 - the job reports any failure as its status + job.error = str(err) or type(err).__name__ + job.touch("failed") + + def finish(task: asyncio.Task[None]) -> None: + # Runs even when the task is cancelled before its body starts. + if task.cancelled(): + job.touch("cancelled") + job.finished = time.perf_counter() + job.done.set() + + job.task = asyncio.create_task(run()) + job.task.add_done_callback(finish) + + +async def crawl_site( + url: str, + limit: int, + depth: int, + include: list[str], + exclude: list[str], + on_page: Callable[[Page], None], +) -> int: + """Discover up to ``limit`` URLs from ``url`` and fetch them locally; returns the total.""" + from .compat_api import page_from_context + from .core.fetcher import Fetcher + from .models.config import CrawlConfig, DocpullConfig + + config = DocpullConfig(url=url, crawl=CrawlConfig(max_pages=limit * 4, max_depth=depth)) + async with Fetcher(config) as fetcher: + discovered = await fetcher.discover() + urls = [item for item in dict.fromkeys([url, *discovered]) if _path_allowed(item, include, exclude)] + urls = urls[:limit] + semaphore = asyncio.Semaphore(CRAWL_CONCURRENCY) + + async def one(item: str) -> None: + async with semaphore: + ctx = await fetcher.fetch_one(item, save=False) + on_page(page_from_context(item, ctx)) + + await asyncio.gather(*(one(item) for item in urls)) + return len(urls) + + +def _path_allowed(url: str, include: list[str], exclude: list[str]) -> bool: + from urllib.parse import urlsplit + + path = urlsplit(url).path or "/" + if include and not any(re.search(pattern, path) for pattern in include): + return False + return not any(re.search(pattern, path) for pattern in exclude) + + +def _patterns(value: Any, name: str) -> list[str]: + if value is None: + return [] + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + raise JobRequestError(f"{name} must be a list of regular expressions") + for pattern in value: + try: + re.compile(pattern) + except re.error as err: + raise JobRequestError(f"Invalid {name} pattern {pattern!r}: {err}") from err + return value + + +def _iso(moment: datetime) -> str: + return moment.isoformat().replace("+00:00", "Z") + + +# Firecrawl /v1/crawl and /v2/crawl ------------------------------------------------ + + +async def create_crawl(app: CompatASGIApp, body: dict[str, Any], base: str) -> tuple[int, dict[str, Any]]: + url = body.get("url") + if not isinstance(url, str) or not url.strip(): + raise JobRequestError("url is required") + limit = body.get("limit", DEFAULT_CRAWL_PAGES) + if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= MAX_CRAWL_PAGES: + raise JobRequestError(f"limit must be an integer from 1 to {MAX_CRAWL_PAGES}") + raw_depth = body.get("maxDiscoveryDepth", body.get("maxDepth", 2)) + if isinstance(raw_depth, bool) or not isinstance(raw_depth, int) or raw_depth < 0: + raise JobRequestError("maxDiscoveryDepth must be a non-negative integer") + include = _patterns(body.get("includePaths"), "includePaths") + exclude = _patterns(body.get("excludePaths"), "excludePaths") + job = app.jobs.add("crawl", body, "") + + async def work(job: Job) -> None: + job.total = await app.crawl(url.strip(), limit, max(1, raw_depth), include, exclude, job.pages.append) + + start(job, work) + return 200, {"success": True, "id": job.id, "url": f"{base}/{job.id}"} + + +def crawl_status(app: CompatASGIApp, job_id: str) -> tuple[int, dict[str, Any]]: + job = app.jobs.get(job_id, "crawl") + if job is None: + return 404, {"success": False, "error": "Crawl job not found"} + status = {"queued": "scraping", "running": "scraping"}.get(job.status, job.status) + fetched = [page for page in job.pages if not page.error] + data = [ + { + "markdown": page.markdown, + "metadata": { + "sourceURL": page.url, + "url": page.url, + "title": page.title, + "statusCode": page.status_code, + }, + } + for page in fetched + ] + payload: dict[str, Any] = { + "success": job.status != "failed", + "status": status, + "total": max(job.total, len(job.pages)), + "completed": len(fetched), + "creditsUsed": 0, + "expiresAt": _iso(job.created_at + JOB_TTL), + "next": None, + "data": data, + } + if job.error: + payload["error"] = job.error + return 200, payload + + +def cancel_crawl(app: CompatASGIApp, job_id: str) -> tuple[int, dict[str, Any]]: + job = app.jobs.get(job_id, "crawl") + if job is None: + return 404, {"success": False, "error": "Crawl job not found"} + if job.task and not job.done.is_set(): + job.task.cancel() + return 200, {"status": "cancelled"} + + +# Shared research work for Parallel task runs and Tavily research --------------- + + +def _research_work( + app: CompatASGIApp, query: str, include: Any = None, exclude: Any = None +) -> Callable[[Job], Awaitable[None]]: + async def work(job: Job) -> None: + response = await app.run_search(query, 5, include, exclude, True) + job.result = await answer_from_search(response.query, response, llm=app.llm) + + return work + + +def _question(value: Any) -> str: + if isinstance(value, str) and value.strip(): + return value.strip() + if isinstance(value, dict) and value: + return " ".join(f"{key}: {item}" for key, item in value.items()) + raise JobRequestError("input must be a non-empty string or object") + + +# Parallel /v1/tasks/runs ------------------------------------------------------- + + +def _task_run(job: Job) -> dict[str, Any]: + return { + "run_id": job.id, + "interaction_id": job.id, + "status": job.status, + "is_active": not job.done.is_set(), + "processor": job.request.get("processor", "docpull"), + "metadata": job.request.get("metadata"), + "error": {"message": job.error} if job.error else None, + "warnings": None, + "created_at": _iso(job.created_at), + "modified_at": _iso(job.modified_at), + } + + +async def create_task_run(app: CompatASGIApp, body: dict[str, Any]) -> tuple[int, dict[str, Any]]: + query = _question(body.get("input")) + app.require_search() + job = app.jobs.add("task", body, "trun_") + start(job, _research_work(app, query)) + return 200, _task_run(job) + + +def task_run(app: CompatASGIApp, job_id: str) -> tuple[int, dict[str, Any]]: + job = app.jobs.get(job_id, "task") + return (200, _task_run(job)) if job else (404, {"error": "Task run not found"}) + + +async def task_run_result(app: CompatASGIApp, job_id: str, timeout: float) -> tuple[int, dict[str, Any]]: + job = app.jobs.get(job_id, "task") + if job is None: + return 404, {"error": "Task run not found"} + try: + await asyncio.wait_for(job.done.wait(), timeout=timeout) + except TimeoutError: + return 408, {"error": "Task run is still active", "run": _task_run(job)} + if job.result is None: + return 422, {"error": job.error or "Task run failed", "run": _task_run(job)} + citations = [ + { + "url": source["url"], + "title": source["title"], + "excerpts": [e.text for e in job.result.evidence if e.source_id == source["id"]], + } + for source in job.result.sources + ] + basis = [ + { + "field": "output", + "citations": citations, + "reasoning": f"Answer built from {len(job.result.evidence)} ranked passages ({job.result.mode}).", + "confidence": None, + } + ] + output = {"type": "text", "content": job.result.answer, "basis": basis} + return 200, {"run": _task_run(job), "output": output} + + +# Tavily /research --------------------------------------------------------------- + + +async def create_research(app: CompatASGIApp, body: dict[str, Any]) -> tuple[int, dict[str, Any]]: + query = _question(body.get("input")) + app.require_search() + job = app.jobs.add("research", body, "") + start(job, _research_work(app, query, body.get("include_domains"), body.get("exclude_domains"))) + return 201, { + "request_id": job.id, + "created_at": _iso(job.created_at), + "status": "pending", + "input": query, + "model": body.get("model", "docpull"), + "response_time": 0, + } + + +def research_status(app: CompatASGIApp, job_id: str) -> tuple[int, dict[str, Any]]: + job = app.jobs.get(job_id, "research") + if job is None: + return 404, {"error": "Research task not found"} + if not job.done.is_set(): + status = "pending" if job.status == "queued" else "in_progress" + return 202, {"request_id": job.id, "status": status, "response_time": job.elapsed} + payload: dict[str, Any] = { + "request_id": job.id, + "created_at": _iso(job.created_at), + "status": job.status, + "response_time": job.elapsed, + } + if job.result is not None: + payload["content"] = job.result.answer + payload["sources"] = [ + {"url": s["url"], "title": s["title"], "favicon": None} for s in job.result.sources + ] + else: + payload["error"] = job.error + return 200, payload diff --git a/tests/test_compat_jobs.py b/tests/test_compat_jobs.py new file mode 100644 index 0000000..550d800 --- /dev/null +++ b/tests/test_compat_jobs.py @@ -0,0 +1,176 @@ +"""Tests for the async crawl, task-run, and research job routes.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from docpull.compat_api import CompatASGIApp, Page +from docpull.compat_jobs import _path_allowed +from docpull.web_search import SearchHit, SearchResponse +from tests.test_compat_api import _post + +DOC = "# Limits\n\nThe API allows 100 requests per minute per key.\n" + + +async def _search(query: str, max_results: int, include: list[str], exclude: list[str], content: bool): + return SearchResponse( + query=query, + backend="fake", + hits=[SearchHit(url="https://docs.example.com/limits", title="Limits", content=DOC)], + ) + + +async def _crawl( + url: str, limit: int, depth: int, include: list[str], exclude: list[str], on_page: Any +) -> int: + urls = [f"{url}/a", f"{url}/b", f"{url}/missing"][:limit] + for item in urls: + await asyncio.sleep(0) + error = "HTTP 404" if item.endswith("missing") else None + on_page( + Page(url=item, markdown="" if error else f"# {item}\n", title=item, status_code=200, error=error) + ) + return len(urls) + + +async def _settle(app: CompatASGIApp) -> None: + for job in list(app.jobs._jobs.values()): + await job.done.wait() + + +@pytest.mark.asyncio +async def test_firecrawl_crawl_job_lifecycle() -> None: + app = CompatASGIApp(crawl=_crawl) + headers = [(b"host", b"127.0.0.1:8765")] + + status, created = await _post( + app, "/firecrawl/v2/crawl", {"url": "https://example.com", "limit": 3}, headers=headers + ) + assert status == 200 and created["success"] is True + assert created["url"] == f"http://127.0.0.1:8765/firecrawl/v2/crawl/{created['id']}" + + await _settle(app) + status, done = await _post(app, f"/v2/crawl/{created['id']}", {}, method="GET") + assert status == 200 + assert done["status"] == "completed" + assert done["total"] == 3 and done["completed"] == 2 + assert done["data"][0]["metadata"]["sourceURL"] == "https://example.com/a" + assert done["creditsUsed"] == 0 + + assert (await _post(app, "/v1/crawl/unknown", {}, method="GET"))[0] == 404 + assert (await _post(app, "/v2/crawl", {"url": "https://example.com", "limit": 9999}))[0] == 400 + assert (await _post(app, "/v2/crawl", {"url": "https://example.com", "includePaths": ["("]}))[0] == 400 + + +@pytest.mark.asyncio +async def test_crawl_cancel_and_failure_status() -> None: + gate = asyncio.Event() + + async def slow(*_args: Any) -> int: + await gate.wait() + return 0 + + app = CompatASGIApp(crawl=slow) + _, created = await _post(app, "/v1/crawl", {"url": "https://example.com"}) + _, running = await _post(app, f"/v1/crawl/{created['id']}", {}, method="GET") + assert running["status"] == "scraping" + _, cancelled = await _post(app, f"/v1/crawl/{created['id']}", {}, method="DELETE") + assert cancelled == {"status": "cancelled"} + await _settle(app) + assert (await _post(app, f"/v1/crawl/{created['id']}", {}, method="GET"))[1]["status"] == "cancelled" + + async def broken(*_args: Any) -> int: + raise RuntimeError("discovery failed") + + failing = CompatASGIApp(crawl=broken) + _, created = await _post(failing, "/v2/crawl", {"url": "https://example.com"}) + await _settle(failing) + _, failed = await _post(failing, f"/v2/crawl/{created['id']}", {}, method="GET") + assert failed["status"] == "failed" and failed["error"] == "discovery failed" + + +@pytest.mark.asyncio +async def test_parallel_task_run_lifecycle() -> None: + app = CompatASGIApp(search=_search) + + status, run = await _post(app, "/v1/tasks/runs", {"input": "requests per minute", "processor": "base"}) + assert status == 200 + assert run["run_id"].startswith("trun_") and run["status"] == "queued" and run["is_active"] is True + assert run["processor"] == "base" + + status, result = await _post(app, f"/v1/tasks/runs/{run['run_id']}/result", {}, method="GET") + assert status == 200 + assert result["run"]["status"] == "completed" and result["run"]["is_active"] is False + assert result["output"]["type"] == "text" and "[S1]" in result["output"]["content"] + citation = result["output"]["basis"][0]["citations"][0] + assert citation["url"] == "https://docs.example.com/limits" + assert "100 requests per minute" in citation["excerpts"][0] + + assert (await _post(app, f"/v1/tasks/runs/{run['run_id']}", {}, method="GET"))[1]["status"] == "completed" + assert (await _post(app, "/v1/tasks/runs", {"input": ""}))[0] == 400 + assert (await _post(CompatASGIApp(), "/v1/tasks/runs", {"input": "q"}))[0] == 501 + + +@pytest.mark.asyncio +async def test_task_run_result_times_out_while_active() -> None: + gate = asyncio.Event() + + async def blocked(*_args: Any) -> SearchResponse: + await gate.wait() + return await _search("q", 5, [], [], True) + + app = CompatASGIApp(search=blocked) + _, run = await _post(app, "/v1/tasks/runs", {"input": "requests per minute"}) + messages: list[Any] = [] + + async def receive() -> dict[str, Any]: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: dict[str, Any]) -> None: + messages.append(message) + + scope = { + "type": "http", + "method": "GET", + "path": f"/v1/tasks/runs/{run['run_id']}/result", + "query_string": b"timeout=0.01", + "headers": [], + } + await app(scope, receive, send) + assert messages[0]["status"] == 408 + gate.set() + await _settle(app) + + +@pytest.mark.asyncio +async def test_tavily_research_lifecycle() -> None: + app = CompatASGIApp(search=_search) + + status, created = await _post(app, "/research", {"input": "requests per minute", "model": "mini"}) + assert status == 201 and created["status"] == "pending" and created["model"] == "mini" + + status, pending = await _post(app, f"/research/{created['request_id']}", {}, method="GET") + assert status == 202 and pending["status"] in {"pending", "in_progress"} + + await _settle(app) + status, done = await _post(app, f"/tavily/research/{created['request_id']}", {}, method="GET") + assert status == 200 and done["status"] == "completed" + assert "[S1]" in done["content"] + assert done["sources"] == [{"url": "https://docs.example.com/limits", "title": "Limits", "favicon": None}] + assert (await _post(app, "/research/nope", {}, method="GET"))[0] == 404 + + +@pytest.mark.asyncio +async def test_job_routes_require_the_api_key() -> None: + app = CompatASGIApp(search=_search, crawl=_crawl, api_key="k") + assert (await _post(app, "/v1/tasks/runs", {"input": "q"}))[0] == 401 + assert (await _post(app, "/research/x", {}, method="GET"))[0] == 401 + + +def test_crawl_path_filters_are_regular_expressions() -> None: + assert _path_allowed("https://e.com/docs/a", [r"^/docs/"], []) + assert not _path_allowed("https://e.com/blog/a", [r"^/docs/"], []) + assert not _path_allowed("https://e.com/docs/private/x", [r"^/docs/"], [r"private"])