diff --git a/src/plyngent/lmproto/deepseek/search/__init__.py b/src/plyngent/lmproto/deepseek/search/__init__.py new file mode 100644 index 0000000..a4b4149 --- /dev/null +++ b/src/plyngent/lmproto/deepseek/search/__init__.py @@ -0,0 +1,8 @@ +from .client import DeepseekSearchClient as DeepseekSearchClient +from .model import DEFAULT_SEARCH_COUNT as DEFAULT_SEARCH_COUNT +from .model import MAX_QUERY_CHARS as MAX_QUERY_CHARS +from .model import MAX_SEARCH_COUNT as MAX_SEARCH_COUNT +from .model import SEARCH_PATH as SEARCH_PATH +from .model import SearchRequest as SearchRequest +from .model import SearchResponse as SearchResponse +from .model import SearchResult as SearchResult diff --git a/src/plyngent/lmproto/deepseek/search/client.py b/src/plyngent/lmproto/deepseek/search/client.py new file mode 100644 index 0000000..63f7717 --- /dev/null +++ b/src/plyngent/lmproto/deepseek/search/client.py @@ -0,0 +1,153 @@ +"""DeepSeek's Search API: one request per query, no model round. + +The endpoint is a plain JSON API on the platform base URL, authenticated with +the same bearer key as the model APIs. It replaces the older server-side search +turn (``POST /messages`` with a hosted ``web_search`` tool), which spent a whole +completion round — and its input tokens — just to pick the query, and could only +return titles and URLs. A Search API response already carries a snippet and a +model-facing summary, so the local ``web_search`` tool can hand the model +something readable on the first call. + +``count`` does not affect billing (one successful request costs the same for 1 +or 50 results); only successful requests are billable, so the retries below are +safe. The account-level limit is 5 QPS and is shared across the Search and Fetch +APIs independently of model concurrency. +""" + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +import msgspec +import niquests +from niquests.auth import BearerTokenAuth + +from plyngent.lmproto.openai_compatible.client import ( + http_error_message, + read_response_body, +) +from plyngent.lmproto.search import SearchHit + +from .model import ( + DEFAULT_SEARCH_COUNT, + MAX_QUERY_CHARS, + MAX_SEARCH_COUNT, + SEARCH_PATH, + SearchRequest, + SearchResponse, + SearchResult, +) + +if TYPE_CHECKING: + from plyngent.lmproto.openai_compatible.config import OpenAIConfig + +_HTTP_ERROR = 400 +_SERVER_ERROR = 500 +# A response is retried when the server (or the account's QPS window) is the +# problem; a request-shaped error is never retried. +_RETRYABLE_STATUS = frozenset({408, 429}) +RETRY_DELAYS: tuple[float, ...] = (0.5, 1.0) +_SNIPPET_FALLBACK_CHARS = 200 + + +class DeepseekSearchClient: + """``POST /v1/tools/search``: one query in, ranked hits out.""" + + session: niquests.AsyncSession + encoder: msgspec.json.Encoder + decoder: msgspec.json.Decoder[SearchResponse] + + def __init__(self, config: OpenAIConfig) -> None: + self.session = niquests.AsyncSession( + base_url=config.base_url, + auth=BearerTokenAuth(config.access_key_or_token), + timeout=config.timeout, + ) + self.encoder = msgspec.json.Encoder() + self.decoder = msgspec.json.Decoder(SearchResponse) + + async def search( + self, + query: str, + *, + count: int = DEFAULT_SEARCH_COUNT, + need_content: bool = False, + ) -> list[SearchHit]: + """Run *query* and return the pages the index ranked. + + The query is forwarded with its search syntax intact (``site:``, + ``lang:``, ``after:``, quoted phrases, …), so the caller decides how + specific to be. ``need_content`` asks the index for page bodies; the + local tool leaves it off and reads a page with ``fetch`` instead. + """ + stripped = query.strip() + if not stripped: + msg = "search query must not be empty" + raise ValueError(msg) + if len(query) > MAX_QUERY_CHARS: + msg = f"search query must be at most {MAX_QUERY_CHARS} characters ({len(query)} given)" + raise ValueError(msg) + payload = SearchRequest( + query=query, + count=max(1, min(count, MAX_SEARCH_COUNT)), + need_content=need_content, + ) + body = await self._post(SEARCH_PATH, self.encoder.encode(payload)) + parsed = self.decoder.decode(body) + return [hit for hit in map(_to_hit, parsed.results) if hit.url] + + async def _post(self, path: str, data: bytes) -> bytes: + """POST *data*, retrying the statuses that mean "wait and try again".""" + for delay in RETRY_DELAYS: + status, body = await self._send(path, data) + if not _retryable(status): + return _require_body(status, body) + await asyncio.sleep(delay) + status, body = await self._send(path, data) + return _require_body(status, body) + + async def _send(self, path: str, data: bytes) -> tuple[int, bytes | str | None]: + resp = await self.session.post( + path, + data=data, + headers={"Content-Type": "application/json"}, + stream=False, + ) + body = await read_response_body(resp) + return int(getattr(resp, "status_code", 0) or 0), body + + +def _retryable(status: int) -> bool: + return status in _RETRYABLE_STATUS or status >= _SERVER_ERROR + + +def _require_body(status: int, body: bytes | str | None) -> bytes: + """Return a successful body; raise the endpoint's error otherwise.""" + if status >= _HTTP_ERROR: + detail = http_error_message(status, body, what="search") or f"search HTTP {status}" + raise RuntimeError(detail) + if body is None: + msg = "search response body is empty" + raise RuntimeError(msg) + return body.encode() if isinstance(body, str) else bytes(body) + + +def _to_hit(result: SearchResult) -> SearchHit: + snippet = _collapse(result.snippet) + if not snippet: + snippet = _shorten(_collapse(result.summary)) + return SearchHit(title=result.title.strip(), url=result.url.strip(), snippet=snippet) + + +def _collapse(text: str) -> str: + return " ".join(text.split()) + + +def _shorten(text: str) -> str: + if len(text) <= _SNIPPET_FALLBACK_CHARS: + return text + return text[: _SNIPPET_FALLBACK_CHARS - 1].rstrip() + "…" + + +__all__ = ["RETRY_DELAYS", "DeepseekSearchClient"] diff --git a/src/plyngent/lmproto/deepseek/search/model.py b/src/plyngent/lmproto/deepseek/search/model.py new file mode 100644 index 0000000..d14fe71 --- /dev/null +++ b/src/plyngent/lmproto/deepseek/search/model.py @@ -0,0 +1,49 @@ +"""Wire models for DeepSeek's Search API (``POST /v1/tools/search``).""" + +from __future__ import annotations + +import msgspec + +SEARCH_PATH = "/v1/tools/search" +DEFAULT_SEARCH_COUNT = 10 +MAX_SEARCH_COUNT = 50 +MAX_QUERY_CHARS = 128 + + +class SearchRequest(msgspec.Struct, omit_defaults=True): + """Body of one search request (unknown optional fields are left out).""" + + query: str + count: int = DEFAULT_SEARCH_COUNT + ori_question: str | None = None + need_content: bool = False + + +class SearchResult(msgspec.Struct, omit_defaults=True): + """One ranked result; unknown fields are ignored on decode.""" + + title: str = "" + url: str = "" + snippet: str = "" + # A 600-800 character summary "intended to be read by a model"; used as the + # snippet when the index has no short excerpt. + summary: str = "" + rerank_score: float = 0.0 + authority: int = -1 + + +class SearchResponse(msgspec.Struct): + """Search results; empty when the query matched nothing (still HTTP 200).""" + + results: list[SearchResult] + + +__all__ = [ + "DEFAULT_SEARCH_COUNT", + "MAX_QUERY_CHARS", + "MAX_SEARCH_COUNT", + "SEARCH_PATH", + "SearchRequest", + "SearchResponse", + "SearchResult", +] diff --git a/src/plyngent/lmproto/search.py b/src/plyngent/lmproto/search.py new file mode 100644 index 0000000..081e9db --- /dev/null +++ b/src/plyngent/lmproto/search.py @@ -0,0 +1,27 @@ +"""Provider-neutral models shared by the local ``web_search`` sources. + +A search source answers a query with a list of pages; the local tool in +``plyngent.tools.net.search`` formats them and the CLI wires whichever sources +``provider.search_sources`` names. Keeping the result type here (rather than in +one provider's package) lets the generic sources — Bing's RSS, DeepSeek's Search +API — speak the same shape without importing each other. +""" + +from __future__ import annotations + +import msgspec + + +class SearchHit(msgspec.Struct, frozen=True, omit_defaults=True): + """One page a search source matched. + + ``snippet`` is whatever short excerpt the source ships: a real snippet from + Bing, or DeepSeek's ``summary`` when its index has no snippet for the page. + """ + + title: str = "" + url: str = "" + snippet: str = "" + + +__all__ = ["SearchHit"] diff --git a/tests/test_lmproto/test_deepseek_search.py b/tests/test_lmproto/test_deepseek_search.py new file mode 100644 index 0000000..2a038e3 --- /dev/null +++ b/tests/test_lmproto/test_deepseek_search.py @@ -0,0 +1,199 @@ +"""DeepSeek's Search API client (``POST /v1/tools/search``).""" + +from __future__ import annotations + +import msgspec +import pytest + +from plyngent.lmproto.deepseek.search import ( + DeepseekSearchClient, + SearchRequest, +) +from plyngent.lmproto.openai_compatible.config import OpenAIConfig +from plyngent.lmproto.search import SearchHit + + +class _Resp: + def __init__(self, status_code: int, body: bytes) -> None: + self.status_code = status_code + self.content = body + + +def _body(*results: dict[str, object]) -> bytes: + return msgspec.json.encode({"results": list(results)}) + + +def _client() -> DeepseekSearchClient: + return DeepseekSearchClient( + OpenAIConfig(access_key_or_token="sk-test", base_url="https://api.deepseek.com"), + ) + + +def _decode_request(data: object) -> dict[str, object]: + assert isinstance(data, bytes) + return msgspec.json.decode(data, type=dict[str, object]) + + +@pytest.mark.asyncio +async def test_search_posts_the_query_and_maps_the_hits(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + sent: dict[str, object] = {} + + async def fake_post(path: str, **kwargs: object) -> _Resp: + sent["path"] = path + sent.update(kwargs) + return _Resp( + 200, + _body( + { + "rerank_score": 2.4, + "title": " DeepSeek V4 released ", + "url": " https://example.com/v4 ", + "snippet": "Tool calling and long context.", + "summary": "A longer summary.", + "authority": 3, + }, + {"title": "Unusable", "url": ""}, + ), + ) + + monkeypatch.setattr(client.session, "post", fake_post) + hits = await client.search("DeepSeek V4", count=5) + + assert hits == [ + SearchHit( + title="DeepSeek V4 released", + url="https://example.com/v4", + snippet="Tool calling and long context.", + ) + ] + assert sent["path"] == "/v1/tools/search" + assert sent["stream"] is False + payload = _decode_request(sent["data"]) + assert payload == {"query": "DeepSeek V4", "count": 5} + + +@pytest.mark.asyncio +async def test_summary_stands_in_when_the_index_has_no_snippet(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + + async def fake_post(path: str, **kwargs: object) -> _Resp: + return _Resp( + 200, + _body( + { + "title": "T", + "url": "https://example.com/1", + "snippet": " ", + "summary": " A 600 character\tsummary\nspread over lines. ", + } + ), + ) + + monkeypatch.setattr(client.session, "post", fake_post) + hits = await client.search("q") + assert hits == [ + SearchHit(title="T", url="https://example.com/1", snippet="A 600 character summary spread over lines.") + ] + + +@pytest.mark.asyncio +async def test_long_summaries_are_cut_for_the_tool_line(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + + async def fake_post(path: str, **kwargs: object) -> _Resp: + return _Resp(200, _body({"title": "T", "url": "https://example.com/1", "summary": "x" * 500})) + + monkeypatch.setattr(client.session, "post", fake_post) + hits = await client.search("q") + assert len(hits[0].snippet) == 200 + assert hits[0].snippet.endswith("…") + + +@pytest.mark.asyncio +async def test_empty_results_are_a_normal_answer(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + + async def fake_post(path: str, **kwargs: object) -> _Resp: + return _Resp(200, _body()) + + monkeypatch.setattr(client.session, "post", fake_post) + assert await client.search("nothing matches this") == [] + + +@pytest.mark.asyncio +async def test_count_is_clamped_to_the_documented_range(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + seen: list[dict[str, object]] = [] + + async def fake_post(path: str, **kwargs: object) -> _Resp: + seen.append(_decode_request(kwargs["data"])) + return _Resp(200, _body()) + + monkeypatch.setattr(client.session, "post", fake_post) + _ = await client.search("q", count=0) + _ = await client.search("q", count=99) + assert [payload["count"] for payload in seen] == [1, 50] + + +@pytest.mark.asyncio +async def test_an_oversized_query_is_refused_before_the_request(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + calls: list[object] = [] + + async def fake_post(path: str, **kwargs: object) -> _Resp: + calls.append(path) + return _Resp(200, _body()) + + monkeypatch.setattr(client.session, "post", fake_post) + with pytest.raises(ValueError, match="at most 128 characters"): + _ = await client.search("x" * 129) + assert calls == [] + + +@pytest.mark.asyncio +async def test_the_endpoint_error_detail_is_raised(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + + async def fake_post(path: str, **kwargs: object) -> _Resp: + return _Resp(400, b'{"error": {"message": "\'count\' must be between 1 and 50"}}') + + monkeypatch.setattr(client.session, "post", fake_post) + with pytest.raises(RuntimeError, match="must be between 1 and 50"): + _ = await client.search("q") + + +@pytest.mark.asyncio +async def test_a_retryable_status_is_tried_again(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + monkeypatch.setattr("plyngent.lmproto.deepseek.search.client.RETRY_DELAYS", (0.0,)) + statuses = [429, 200] + + async def fake_post(path: str, **kwargs: object) -> _Resp: + return _Resp(statuses.pop(0), _body({"title": "T", "url": "https://example.com/1"})) + + monkeypatch.setattr(client.session, "post", fake_post) + hits = await client.search("q") + assert [hit.url for hit in hits] == ["https://example.com/1"] + assert statuses == [] + + +@pytest.mark.asyncio +async def test_a_balance_error_is_not_retried(monkeypatch: pytest.MonkeyPatch) -> None: + client = _client() + monkeypatch.setattr("plyngent.lmproto.deepseek.search.client.RETRY_DELAYS", (0.0,)) + calls: list[int] = [] + + async def fake_post(path: str, **kwargs: object) -> _Resp: + calls.append(1) + return _Resp(402, b'{"error": {"message": "Insufficient Balance"}}') + + monkeypatch.setattr(client.session, "post", fake_post) + with pytest.raises(RuntimeError, match="Insufficient Balance"): + _ = await client.search("q") + assert calls == [1] + + +def test_the_request_model_omits_what_was_not_set() -> None: + encoded = msgspec.json.encode(SearchRequest(query="q")) + assert msgspec.json.decode(encoded, type=dict[str, object]) == {"query": "q"}