diff --git a/cookbook/agentic_retrieval.ipynb b/cookbook/agentic_retrieval.ipynb index dacaf01ea..518235d72 100644 --- a/cookbook/agentic_retrieval.ipynb +++ b/cookbook/agentic_retrieval.ipynb @@ -339,11 +339,7 @@ "source": [ "query = \"What are the evaluation methods used in this paper?\"\n", "\n", - "for chunk in pi_client.chat_completions(\n", - " messages=[{\"role\": \"user\", \"content\": query}],\n", - " doc_id=doc_id,\n", - " stream=True\n", - "):\n", + "for chunk in pi_client.chat(query, doc_id=doc_id, stream=True):\n", " print(chunk, end='', flush=True)" ] }, @@ -426,11 +422,8 @@ "\n", "full_response = \"\"\n", "\n", - "for chunk in pi_client.chat_completions(\n", - " messages=[{\"role\": \"user\", \"content\": retrieval_prompt}],\n", - " doc_id=doc_id,\n", - " stream=True\n", - "):\n", + "for chunk in pi_client.chat(retrieval_prompt, doc_id=doc_id, stream=True,\n", + " show_process=False):\n", " print(chunk, end='', flush=True)\n", " full_response += chunk" ] diff --git a/pageindex/client.py b/pageindex/client.py index 8524bd5f7..cdfc17eea 100644 --- a/pageindex/client.py +++ b/pageindex/client.py @@ -519,8 +519,9 @@ def _local_chat(self) -> bool: return model is not None def _require_own_chat(self, lane: str) -> None: - # The one refusal for chat(protocol=...), the doors behind it, and - # instructions: shared, so the doors cannot drift from chat(). + # The one refusal for the Responses / Messages lanes, the doors + # behind them, and instructions: shared, so the doors cannot drift + # from chat(). if self._local_chat: return if not getattr(self, "api_key", None): @@ -531,7 +532,8 @@ def _require_own_chat(self, lane: str) -> None: raise PageIndexAPIError( f"{lane} drives your own chat model — construct the client " "with chat_model=... (or a chat= model); the managed cloud chat " - "serves the answer lane and chat_completions() only.") + "serves the answer lane and chat(protocol=\"chat_completions\") " + "only.") if not TYPE_CHECKING: # The protocol doors live behind chat(protocol=...); their old @@ -755,11 +757,11 @@ def submit_query(self, doc_id: str, query: str, thinking: bool = False) -> dict[ Cloud-only: the cloud API marks this endpoint deprecated in favor of chat completions, so local mode does not implement it — raises - PageIndexAPIError. Use ``chat_completions`` instead. + PageIndexAPIError. Use ``chat()`` instead. """ return self._require_cloud( "submit_query is cloud-only — the retrieval API is deprecated in " - "favor of chat completions; use chat_completions instead." + "favor of chat completions; use chat() instead." ).submit_query(doc_id=doc_id, query=query, thinking=thinking) def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: @@ -768,11 +770,11 @@ def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: Cloud-only: the cloud API marks this endpoint deprecated in favor of chat completions, so local mode does not implement it — raises - PageIndexAPIError. Use ``chat_completions`` instead. + PageIndexAPIError. Use ``chat()`` instead. """ return self._require_cloud( "get_retrieval is cloud-only — the retrieval API is deprecated in " - "favor of chat completions; use chat_completions instead." + "favor of chat completions; use chat() instead." ).get_retrieval(retrieval_id=retrieval_id) # ---------- CHAT ---------- @@ -784,12 +786,12 @@ def get_retrieval(self, retrieval_id: str) -> dict[str, Any]: def chat( self, messages: Union[str, list[dict[str, Any]]], + *, doc_id: Optional[Union[str, list[str]]] = None, stream: Literal[False] = False, model: Optional[str] = None, reasoning_effort: Optional[str] = None, show_process: Union[bool, Mapping[str, Any], None] = None, - *, protocol: None = None, instructions: Optional[Union[str, list[dict[str, Any]]]] = None, max_turns: Optional[int] = None, @@ -802,8 +804,8 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], - doc_id: Optional[Union[str, list[str]]] = None, *, + doc_id: Optional[Union[str, list[str]]] = None, stream: Literal[True], model: Optional[str] = None, reasoning_effort: Optional[str] = None, @@ -820,13 +822,13 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], + *, doc_id: Optional[Union[str, list[str]]] = None, stream: Literal[False] = False, model: Optional[str] = None, reasoning_effort: Optional[str] = None, show_process: Union[bool, Mapping[str, Any], None] = None, - *, - protocol: Literal["responses", "messages"], + protocol: Literal["chat_completions", "responses", "messages"], instructions: Optional[Union[str, list[dict[str, Any]]]] = None, max_turns: Optional[int] = None, backend: Optional[dict[str, Any]] = None, @@ -838,13 +840,13 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], - doc_id: Optional[Union[str, list[str]]] = None, *, + doc_id: Optional[Union[str, list[str]]] = None, stream: Literal[True], model: Optional[str] = None, reasoning_effort: Optional[str] = None, show_process: Union[bool, Mapping[str, Any], None] = None, - protocol: Literal["responses"], + protocol: Literal["chat_completions", "responses"], instructions: Optional[Union[str, list[dict[str, Any]]]] = None, max_turns: Optional[int] = None, backend: Optional[dict[str, Any]] = None, @@ -856,8 +858,8 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], - doc_id: Optional[Union[str, list[str]]] = None, *, + doc_id: Optional[Union[str, list[str]]] = None, stream: Literal[True], model: Optional[str] = None, reasoning_effort: Optional[str] = None, @@ -874,12 +876,12 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], + *, doc_id: Optional[Union[str, list[str]]] = None, stream: bool = False, model: Optional[str] = None, reasoning_effort: Optional[str] = None, show_process: Union[bool, Mapping[str, Any], None] = None, - *, protocol: None = None, instructions: Optional[Union[str, list[dict[str, Any]]]] = None, max_turns: Optional[int] = None, @@ -892,13 +894,14 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], + *, doc_id: Optional[Union[str, list[str]]] = None, stream: bool = False, model: Optional[str] = None, reasoning_effort: Optional[str] = None, show_process: Union[bool, Mapping[str, Any], None] = None, - *, - protocol: Optional[str] = None, + protocol: Optional[Literal["chat_completions", "responses", + "messages"]] = None, instructions: Optional[Union[str, list[dict[str, Any]]]] = None, max_turns: Optional[int] = None, backend: Optional[dict[str, Any]] = None, @@ -909,13 +912,14 @@ def chat( def chat( self, messages: Union[str, list[dict[str, Any]]], + *, doc_id: Optional[Union[str, list[str]]] = None, stream: bool = False, model: Optional[str] = None, reasoning_effort: Optional[str] = None, show_process: Union[bool, Mapping[str, Any], None] = None, - *, - protocol: Optional[str] = None, + protocol: Optional[Literal["chat_completions", "responses", + "messages"]] = None, instructions: Optional[Union[str, list[dict[str, Any]]]] = None, max_turns: Optional[int] = None, backend: Optional[dict[str, Any]] = None, @@ -933,10 +937,14 @@ def chat( join a stream into one only with ``show_process=False``) and pass it back. - The protocol lanes (``protocol="responses"`` / ``"messages"``): - own-model chat driven natively over the OpenAI Responses API or - Anthropic's Messages API. Input and output are that protocol's own - shapes — the history may carry its transcript (Responses items, or + The protocol lanes: ``protocol="chat_completions"`` is the answer + lane's own engine with its envelope kept — the Chat Completions + response (``choices``/``usage``), or its chunk dicts when + streaming; it is the one protocol the managed cloud chat serves + too. ``protocol="responses"`` / ``"messages"``: own-model chat + driven natively over the OpenAI Responses API or Anthropic's + Messages API. Input and output are that protocol's own shapes — + the history may carry its transcript (Responses items, or Messages content blocks with prior tool_use/tool_result round-trips), and the return is its response envelope, streaming its native events. A round-tripped transcript continues the @@ -947,11 +955,15 @@ def chat( Args: messages: A question string, or the conversation history — - role/content messages on every lane. ``system`` rows join - the managed prompt on the answer lane only, wherever they - sit; the protocol lanes pass rows to the wire as they are - (use ``instructions`` for persona there). With a protocol, - also that protocol's transcript items or content blocks. + role/content messages on every lane. With your own chat + model, ``system`` rows join the managed prompt on the + answer lane and ``protocol="chat_completions"``, wherever + they sit (the managed endpoint forwards them verbatim); + the other protocol lanes pass rows to the wire as they + are (use ``instructions`` for persona there). Responses + and Messages also accept their native transcript items + or content blocks; own-model Chat Completions takes text + history only. doc_id: Document ID or list of IDs to scope the conversation. Keep it identical across a conversation's calls. Local documents: also enforced at the tool layer, not just @@ -997,18 +1009,22 @@ def chat( models expose none on the chat protocol). The labels are not a parse format, and a process stream must not be appended back as conversation history — for the - machine-readable process use ``.events``, or a protocol - lane's transcript. A protocol lane returns that - transcript itself, so ``show_process`` is an error there. - protocol: ``None`` for the answer lane, or ``"responses"`` / - ``"messages"`` — the wire protocol, engine, and - input/output shapes of this call. Own-model chat only. + machine-readable process use ``.events``, or the + Responses / Messages lane's transcript. A protocol lane + returns its own shape, so ``show_process`` is an error + there. + protocol: ``None`` for the answer lane, or + ``"chat_completions"`` / ``"responses"`` / ``"messages"`` + — the wire protocol, engine, and input/output shapes of + this call. Own-model chat only, except + ``"chat_completions"``, which the managed chat serves too. instructions: Own-model chat only — persona or extra guidance appended after the managed system prompt (which stays: it carries the tool guidance and the document context). A string on every lane; with ``protocol="messages"`` also - a list of Messages system blocks. On the answer lane it - precedes any ``system`` rows in the history. + a list of Messages system blocks. On the answer lane and + ``protocol="chat_completions"`` it precedes any ``system`` + rows in the history. max_turns: Own-model chat only — cap on agent turns per call (default 10). The OpenAI lanes raise at the cap; ``protocol="messages"`` returns the truncated run @@ -1017,24 +1033,29 @@ def chat( backend: Own-model chat only — connection overrides for this call's backend, merged over the client's ``chat_backend`` (per-call keys win): LiteLLM's connection params on the - answer lane, the openai / anthropic SDK's client params - on the protocol lanes. Passed through verbatim. + answer lane and ``protocol="chat_completions"``; the + openai / anthropic SDK's client params on Responses / + Messages. Passed through verbatim. extra_headers: Own-model chat only — extra HTTP headers merged into each backend request; caller headers win. LiteLLM's anthropic adapter owns ``anthropic-beta`` on the - answer lane — Anthropic beta flags ride - ``protocol="messages"``. - extra_body: Own-model chat only — the provider's own request - fields beyond this method's parameters, in the lane's - wire names (Responses ``max_output_tokens``, Messages - ``thinking`` / ``top_k``), merged last so they win. - Answer lane: LiteLLM's own params, mapped or refused per - provider (``response_format`` has no door on - LiteLLM-routed models); protocol lanes: verbatim into - the request body. The managed prompt, conversation and - tools are not fields here (``system`` / ``instructions`` - / ``input`` / ``messages`` / ``tools`` are refused); - extend the prompt with ``instructions=``. Credentials + answer lane and ``protocol="chat_completions"`` — + Anthropic beta flags ride ``protocol="messages"``. + extra_body: The wire's own request fields beyond this + method's parameters, in the lane's wire names (Responses + ``max_output_tokens``, Messages ``thinking`` / ``top_k``; + the managed chat endpoint's ``temperature`` / + ``enable_citations``), merged last so they win. + The managed endpoint, Responses / Messages, and + OpenAI-compatible chat backends take these verbatim in + the request body. Other own-model chat backends take + LiteLLM's own params, mapped or refused per provider + (``response_format`` is unsupported there). The managed + prompt, conversation and tools are not fields here (``system`` / + ``instructions`` / ``input`` / ``messages`` / ``tools`` + are refused); extend the prompt with ``instructions=`` or a + leading system row in ``messages``. ``stream`` / ``doc_id`` + are refused too: each has its own argument. Credentials belong in ``backend``, never here. Returns: @@ -1047,24 +1068,27 @@ def chat( ``{"type": "tool_call", "call_id", "name", "arguments"}``, ``{"type": "tool_result", "call_id", "name", "output"}`` - protocol lane, stream=False: the protocol's response - envelope — Responses: ``output`` plus an ``items`` + envelope — Chat Completions: ``choices`` and ``usage``; + Responses: ``output`` plus an ``items`` transcript and cross-turn ``usage``; Messages: the final message with a ``messages`` turn sequence and aggregated ``usage`` - protocol lane, stream=True: an iterator of the protocol's own stream events """ - if protocol not in (None, "responses", "messages"): + if protocol not in (None, "chat_completions", "responses", + "messages"): raise PageIndexAPIError( - "protocol selects the wire: \"responses\" (OpenAI Responses) " - "or \"messages\" (Anthropic Messages), or leave it unset for " + "protocol selects the wire: \"chat_completions\" (OpenAI Chat " + "Completions), \"responses\" (OpenAI Responses) or " + "\"messages\" (Anthropic Messages), or leave it unset for " f"the answer lane — got {protocol!r}.") if (protocol is not None and show_process is not False and show_process is not None): raise PageIndexAPIError( "show_process weaves the answer lane's run; with " - f"protocol={protocol!r} the run comes back as the protocol's " - "own transcript and events — drop show_process, or drop " + f"protocol={protocol!r} the return is the protocol's own " + "envelope and events — drop show_process, or drop " "protocol for the woven text stream.") if show_process is not False and show_process is not None: from .local_chat import _process_options @@ -1079,7 +1103,9 @@ def chat( "instructions blocks are the Messages protocol's shape — " "with protocol=\"messages\" they append after the managed " "system blocks; the other lanes take a string.") - if protocol is not None: + from .local_chat import _refuse_skeleton + _refuse_skeleton(extra_body) + if protocol in ("responses", "messages"): self._require_own_chat(f"chat(protocol={protocol!r})") if protocol == "responses": body = extra_body @@ -1125,6 +1151,12 @@ def chat( # then the history's own system rows. messages = [{"role": "system", "content": instructions}, *messages] + if protocol == "chat_completions": + return self.chat_completions( + messages, stream=stream, stream_metadata=True, doc_id=doc_id, + model=model, max_turns=max_turns, + reasoning_effort=reasoning_effort, extra_body=extra_body, + extra_headers=extra_headers, backend=backend) if stream: # the default means "on where available" resolved = True if show_process is None else show_process @@ -1179,7 +1211,12 @@ def chat_completions( backend: Optional[dict[str, Any]] = None, ) -> Union[dict[str, Any], Iterator[str], Iterator[dict[str, Any]]]: """ - PageIndex Chat Completions: document QA in one call. + Kept for existing code — new code calls ``chat()``. Everything + here is ``chat(protocol="chat_completions")``: the same engine + and envelope, with this method's sampling fields riding + ``extra_body`` under their wire names. The one exception is the + text-only stream (``stream=True`` without ``stream_metadata``): + that is ``chat(stream=True, show_process=False)``. With no chat model configured (a plain cloud client): the managed hosted chat endpoint. With one — local mode, or a cloud client @@ -1238,14 +1275,15 @@ def chat_completions( its own thinking control, and the values mean what the backend says they mean. Unset sends nothing (the backend's default applies). - extra_body: Own-model chat only — extra request fields beyond this - method's parameters, merged last so they win. - OpenAI-compatible backends take them verbatim in the + extra_body: Extra request fields beyond this method's + parameters, merged last so they win. The managed endpoint + and OpenAI-compatible backends take them verbatim in the request body; LiteLLM-routed providers take them as LiteLLM's own params (mapped or refused per provider). The managed prompt, conversation and tools are not fields here (``system`` / ``instructions`` / ``input`` / - ``messages`` / ``tools`` are refused). Credentials belong + ``messages`` / ``tools`` are refused), nor are ``stream`` + / ``doc_id``: each has its own argument. Credentials belong in ``backend``, never here. extra_headers: Own-model chat only — extra HTTP headers merged into each backend request; caller headers win. One exception: @@ -1270,6 +1308,8 @@ def chat_completions( "messages must be a non-empty string or a list of " "message dicts.") messages = [{"role": "user", "content": messages}] + from .local_chat import _refuse_skeleton + _refuse_skeleton(extra_body) if self._local_chat: from .local_chat import run_chat_completions return run_chat_completions( @@ -1287,10 +1327,10 @@ def chat_completions( "chat_model=... to run the agent with your own model.") if (model or max_turns is not None or top_p is not None or max_tokens is not None or reasoning_effort - or extra_body or extra_headers or backend): + or extra_headers or backend): raise PageIndexAPIError( "model, max_turns, top_p, max_tokens, reasoning_effort, " - "extra_body, extra_headers and backend drive your own chat " + "extra_headers and backend drive your own chat " "model, which this client does not configure — construct " "the client with chat_model=... (or a chat= model) to run the " "agent in your process, or drop them to use the managed " @@ -1300,7 +1340,7 @@ def chat_completions( return cast(CloudAPI, self._api).chat_completions( messages=messages, stream=stream, doc_id=doc_id, temperature=temperature, stream_metadata=stream_metadata, - enable_citations=enable_citations, + enable_citations=enable_citations, extra_body=extra_body, ) def _responses( @@ -1336,10 +1376,11 @@ def _responses( Requires a backend that supports the Responses API; backends that only speak chat.completions should use - ``chat_completions()``. Provider-prefixed models (``anthropic/…``) - route through LiteLLM's chat.completions adapter and are therefore - refused here — use ``chat_completions()`` or - ``chat(protocol="messages")`` for those. + ``chat(protocol="chat_completions")``. Provider-prefixed models + (``anthropic/…``) route through LiteLLM's chat.completions adapter + and are therefore refused here — use + ``chat(protocol="chat_completions")`` or ``chat(protocol="messages")`` + for those. Args: input: A user message string, or a list of Responses input items @@ -1921,7 +1962,7 @@ def agent_instructions( SDK release. Raises PageIndexAPIError if the server cannot be reached. Local: the built-in guidance for the in-process tools. - With ``doc_id`` (str or list, same shape as ``chat_completions``), + With ``doc_id`` (str or list, same shape as ``chat``), appends the target documents' names and metadata and directs the agent to work within them. Raises PageIndexAPIError if a doc_id does not exist, or if its name is shadowed by a newer same-name diff --git a/pageindex/cloud_api.py b/pageindex/cloud_api.py index a676cf3f0..a25f72606 100644 --- a/pageindex/cloud_api.py +++ b/pageindex/cloud_api.py @@ -197,7 +197,8 @@ def chat_completions( doc_id: Optional[Union[str, List[str]]] = None, temperature: Optional[float] = None, stream_metadata: bool = False, - enable_citations: bool = False + enable_citations: bool = False, + extra_body: Optional[Dict[str, Any]] = None, ) -> Union[Dict[str, Any], Iterator[str], Iterator[Dict[str, Any]]]: """ PageIndex Chat Completions. Optionally scoped to specific PageIndex documents. @@ -209,6 +210,7 @@ def chat_completions( temperature (Optional[float], optional): Sampling temperature. Default is None (uses API default). stream_metadata (bool, optional): If True and stream=True, return raw chunks with metadata instead of just text. Default is False. enable_citations (bool, optional): Enable citation instructions in responses. Default is False. + extra_body (Optional[Dict[str, Any]], optional): Extra request fields, merged into the payload last. Returns: Union[Dict[str, Any], Iterator[str], Iterator[Dict[str, Any]]]: @@ -230,6 +232,8 @@ def chat_completions( if enable_citations: payload["enable_citations"] = enable_citations + payload.update(extra_body or {}) + response = requests.post( f"{self.BASE_URL}/chat/completions/", headers=self._headers(), diff --git a/pageindex/local_chat.py b/pageindex/local_chat.py index e91b90138..16bd57f46 100644 --- a/pageindex/local_chat.py +++ b/pageindex/local_chat.py @@ -62,8 +62,8 @@ def _system_text(content: Any) -> str: def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]": """Validate the chat_completions surface's messages: system/developer content joins the managed instructions; user/assistant history passes - through. Tool-history round-trips belong to the protocol lanes, - chat(protocol=...).""" + through. Tool-history round-trips belong to chat(protocol="responses") + or chat(protocol="messages").""" if not isinstance(messages, list) or not messages: raise PageIndexAPIError("messages must be a non-empty list.") system_texts: list[str] = [] @@ -80,13 +80,15 @@ def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]": if not isinstance(content, str): raise PageIndexAPIError( "content must be a string on this lane; for " - "structured items use chat(protocol=...)." + "structured items use chat(protocol=\"responses\") " + "or chat(protocol=\"messages\")." ) history.append({"role": role, "content": content}) else: raise PageIndexAPIError( - f"Unsupported role {role!r} on this lane. Tool " - "history round-trips belong to chat(protocol=...)." + f"Unsupported role {role!r} on this lane. Tool history " + "round-trips belong to chat(protocol=\"responses\") or " + "chat(protocol=\"messages\")." ) if not history: raise PageIndexAPIError("messages must contain a user or assistant " @@ -199,8 +201,9 @@ def _openai_model(protocol: str, model_name: str, backend=None): f"protocol='responses' cannot drive '{model_name}': " "provider-prefixed models route through LiteLLM, which speaks " "chat.completions, not the Responses API. Use chat() without " - "protocol, or chat_completions() (or protocol='messages' for " - "Anthropic models), or point OPENAI_BASE_URL at a " + "protocol, or protocol='chat_completions' (or " + "protocol='messages' for Anthropic models), or point " + "OPENAI_BASE_URL at a " "Responses-capable backend and use a bare or " "'openai/'-prefixed model name." ) @@ -310,17 +313,32 @@ def _merged_backend(client, backend): _SKELETON_KEYS = frozenset({"system", "instructions", "input", "messages", "tools"}) +_ARGUMENT_KEYS = frozenset({"stream", "doc_id"}) def _refuse_skeleton(extra_body) -> None: """The managed prompt, conversation and tools are the SDK's on every - lane; extra_body merges last, so a caller's copy would replace them.""" - hit = sorted(_SKELETON_KEYS.intersection(extra_body or ())) + lane; extra_body merges last, so a caller's copy would replace them. + Fields with their own argument select the SDK's parser and scope, so + they are refused here too.""" + if extra_body is None: + return + if not isinstance(extra_body, Mapping): + raise PageIndexAPIError( + "extra_body must be a dict of request fields, got " + f"{type(extra_body).__name__}.") + hit = sorted(_SKELETON_KEYS.intersection(extra_body)) if hit: raise PageIndexAPIError( f"extra_body cannot carry {', '.join(hit)}: the managed prompt, " "conversation and tools are the SDK's. Extend the prompt with " - "instructions=; the conversation is the first argument.") + "instructions= or a leading system row; pass the conversation " + "as messages, the first argument.") + hit = sorted(_ARGUMENT_KEYS.intersection(extra_body)) + if hit: + raise PageIndexAPIError( + f"extra_body cannot carry {', '.join(hit)}: use " + f"{' / '.join(key + '=' for key in hit)} instead.") def _openai_agent(client, protocol: str, model_name: str, instructions: str, @@ -944,7 +962,7 @@ def run_chat_completions(client, messages, stream: bool = False, if getattr(client, "api_key", None) else "local mode does not store the block-level OCR data " "citations need.")) - _require_openai_agents("chat_completions") + _require_openai_agents("chat") _validate_max_turns(max_turns) agent, items, model_name = _chat_agent( client, messages, doc_id, model, temperature=temperature, diff --git a/tests/test_client.py b/tests/test_client.py index f6c92f3a4..51eefcda9 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1354,9 +1354,9 @@ def test_folders_are_cloud_only(local_client): # ── local: retrieval endpoints are cloud-only ── def test_retrieval_endpoints_cloud_only(local_client): - with pytest.raises(PageIndexAPIError, match="use chat_completions"): + with pytest.raises(PageIndexAPIError, match=r"use chat\(\)"): local_client.submit_query("any", "q") - with pytest.raises(PageIndexAPIError, match="use chat_completions"): + with pytest.raises(PageIndexAPIError, match=r"use chat\(\)"): local_client.get_retrieval("any") @@ -1537,6 +1537,71 @@ def test_cloud_chat_accepts_query_string(cloud): client.chat_completions(" ") +def test_cloud_chat_extra_body_merges_into_payload(cloud): + client, calls, fake = cloud + fake.payload = {"choices": [{"message": {"content": "ok"}}]} + client.chat("q", doc_id="pi-1", + extra_body={"temperature": 0.2, "enable_citations": True, + "service_tier": "auto"}) + assert calls[-1]["json"] == { + "messages": [{"role": "user", "content": "q"}], "stream": False, + "doc_id": "pi-1", "temperature": 0.2, "enable_citations": True, + "service_tier": "auto"} + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("extra_stream", [False, True]) +@pytest.mark.parametrize("method, options", [ + ("chat", {}), + ("chat", {"protocol": "chat_completions"}), + ("chat_completions", {}), + ("chat_completions", {"stream_metadata": True}), +]) +def test_cloud_chat_rejects_extra_body_stream_before_request( + cloud, stream, extra_stream, method, options): + client, calls, fake = cloud + fake.payload = {"choices": [{"message": {"content": "ok"}}]} + with pytest.raises(PageIndexAPIError, + match=r"extra_body cannot carry stream.*stream="): + getattr(client, method)("q", stream=stream, + extra_body={"stream": extra_stream}, + **options) + assert calls == [] + + +@pytest.mark.parametrize("bad", [["ab"], "messages", 5, [("a", 1)]]) +@pytest.mark.parametrize("method", ["chat", "chat_completions"]) +def test_cloud_chat_rejects_non_dict_extra_body_before_request( + cloud, bad, method): + client, calls, fake = cloud + fake.payload = {"choices": [{"message": {"content": "ok"}}]} + with pytest.raises(PageIndexAPIError, match="extra_body must be a dict"): + getattr(client, method)("q", extra_body=bad) + assert calls == [] + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("extra_doc_id", [ + None, "pi-1", "pi-other", ["pi-1", "pi-other"], +]) +@pytest.mark.parametrize("method, options", [ + ("chat", {}), + ("chat", {"protocol": "chat_completions"}), + ("chat_completions", {}), + ("chat_completions", {"stream_metadata": True}), +]) +def test_cloud_chat_rejects_extra_body_doc_id_before_request( + cloud, stream, extra_doc_id, method, options): + client, calls, fake = cloud + fake.payload = {"choices": [{"message": {"content": "ok"}}]} + with pytest.raises(PageIndexAPIError, + match=r"extra_body cannot carry doc_id.*doc_id="): + getattr(client, method)("q", doc_id="pi-1", stream=stream, + extra_body={"doc_id": extra_doc_id}, + **options) + assert calls == [] + + def test_parse_pages_overlap_counts_union(): from pageindex.client import _parse_pages pages = _parse_pages("1-5000,2000-9000") diff --git a/tests/test_local_chat.py b/tests/test_local_chat.py index 7179c711c..0f17730c3 100644 --- a/tests/test_local_chat.py +++ b/tests/test_local_chat.py @@ -1,6 +1,7 @@ """Local chat surfaces: three protocols over fake backends — no network, no LLM keys. Tool execution runs for real against a seeded local store.""" import asyncio +import inspect import json import sys import types @@ -312,9 +313,6 @@ def test_cloud_guards(monkeypatch): with pytest.raises(PageIndexAPIError, match="own chat model"): cloud.chat_completions([{"role": "user", "content": "x"}], reasoning_effort="low") - with pytest.raises(PageIndexAPIError, match="own chat model"): - cloud.chat_completions([{"role": "user", "content": "x"}], - extra_body={"service_tier": "auto"}) with pytest.raises(PageIndexAPIError, match="own chat model"): cloud.chat_completions([{"role": "user", "content": "x"}], top_p=0.9) with pytest.raises(PageIndexAPIError, match="own chat model"): @@ -3216,6 +3214,70 @@ def test_old_door_names_point_at_chat_protocol(client): client.no_such_thing +def test_chat_protocol_chat_completions_is_the_door(client, monkeypatch): + seen = [] + monkeypatch.setattr(local_chat, "run_chat_completions", + lambda c, messages, **kw: seen.append((messages, kw)) + or "door") + knobs = dict(doc_id="pi-a", model="gpt-x", max_turns=3, + reasoning_effort="low", backend={"api_key": "k"}, + extra_headers={"x": "1"}, extra_body={"seed": 1}) + for streaming in (False, True): + assert client.chat("q", protocol="chat_completions", stream=streaming, + **knobs) == "door" + # the protocol's own stream is its chunk dicts, never text pieces + assert client.chat_completions("q", stream=streaming, + stream_metadata=True, + **knobs) == "door" + assert seen[-2] == seen[-1] + assert seen[0][0] == [{"role": "user", "content": "q"}] + client.chat("q", protocol="chat_completions", instructions="be brief") + assert seen[-1][0] == [{"role": "system", "content": "be brief"}, + {"role": "user", "content": "q"}] + with pytest.raises(PageIndexAPIError, match="show_process"): + client.chat("q", protocol="chat_completions", stream=True, + show_process=True) + assert client.chat_completions("q") == "door" + + +def test_chat_protocol_chat_completions_serves_managed_cloud(monkeypatch): + """Unlike the own-model protocols, the managed cloud chat speaks + chat.completions itself, so the lane opens without a chat model; + the own-model knobs still refuse there.""" + from pageindex import PageIndexClient + cloud = PageIndexClient(api_key="pi-k") + seen = [] + monkeypatch.setattr(cloud._api, "chat_completions", + lambda **kw: seen.append(kw) or {"choices": []}) + assert cloud.chat("q", protocol="chat_completions") == {"choices": []} + assert seen[-1] == {"messages": [{"role": "user", "content": "q"}], + "stream": False, "doc_id": None, "temperature": None, + "stream_metadata": True, "enable_citations": False, + "extra_body": None} + cloud.chat("q", protocol="chat_completions", + extra_body={"temperature": 0.2, "enable_citations": True}) + assert seen[-1]["extra_body"] == {"temperature": 0.2, + "enable_citations": True} + with pytest.raises(PageIndexAPIError, match="extra_body cannot carry"): + cloud.chat("q", protocol="chat_completions", + extra_body={"messages": []}) + with pytest.raises(PageIndexAPIError, match="chat_model="): + cloud.chat("q", protocol="chat_completions", model="m") + with pytest.raises(PageIndexAPIError, match="chat_model="): + cloud.chat("q", protocol="chat_completions", instructions="x") + + +def test_chat_takes_only_messages_by_position(): + """chat_completions() puts stream before doc_id; chat() the reverse. + A positional rewrite must fail loudly, never bind doc_id=True.""" + params = list(inspect.signature(PageIndexClient.chat).parameters.values()) + assert [p.name for p in params[:2]] == ["self", "messages"] + assert {p.kind for p in params[2:]} == {inspect.Parameter.KEYWORD_ONLY} + cloud = PageIndexClient(api_key="pi-k") + with pytest.raises(TypeError, match="positional"): + cloud.chat("q", True, "pi-1") + + def test_chat_protocol_responses_is_the_door(client, monkeypatch): seen = [] monkeypatch.setattr(local_chat, "run_responses", @@ -3297,6 +3359,46 @@ def test_extra_body_refuses_skeleton_keys(): None, None, extra_body={"input": "x"}) +def test_extra_body_refuses_non_dicts_and_argument_keys(): + """The same gate: a non-dict would be splatted into the payload as + fabricated fields; stream / doc_id select the SDK's parser and scope, + so they ride their own arguments on every lane.""" + for bad in (["ab"], "messages", 5, [("a", 1)]): + with pytest.raises(PageIndexAPIError, + match="extra_body must be a dict"): + local_chat._refuse_skeleton(bad) + for key in ("stream", "doc_id"): + with pytest.raises(PageIndexAPIError, + match=rf"extra_body cannot carry {key}: use {key}="): + local_chat._refuse_skeleton({key: True}) + local_chat._refuse_skeleton(None) + local_chat._refuse_skeleton({}) + local_chat._refuse_skeleton({"service_tier": "auto"}) + + +def test_chat_refuses_bad_extra_body_before_any_lane(client, monkeypatch): + """chat() and chat_completions() check extra_body before entering a + lane, so no lane does I/O (or, on Responses, an effort merge) on a + bad value.""" + entered = [] + for door in ("run_chat_completions", "run_responses", "run_messages"): + monkeypatch.setattr(local_chat, door, + lambda c, *a, **kw: entered.append(1)) + for protocol, knobs in ((None, {}), ("chat_completions", {}), + ("responses", {}), + ("messages", {"model": "claude-x"})): + with pytest.raises(PageIndexAPIError, + match="extra_body must be a dict"): + client.chat("q", protocol=protocol, reasoning_effort="low", + extra_body=["ab"], **knobs) + with pytest.raises(PageIndexAPIError, match="cannot carry stream"): + client.chat("q", protocol=protocol, extra_body={"stream": True}, + **knobs) + with pytest.raises(PageIndexAPIError, match="cannot carry stream"): + client.chat_completions("q", extra_body={"stream": True}) + assert entered == [] + + @needs_anthropic def test_messages_extra_body_refuses_skeleton_before_transport(client, monkeypatch):