diff --git a/docs/LOCAL_MODEL_SERVING.md b/docs/LOCAL_MODEL_SERVING.md index d6b5eee..b8f56ad 100644 --- a/docs/LOCAL_MODEL_SERVING.md +++ b/docs/LOCAL_MODEL_SERVING.md @@ -276,3 +276,33 @@ biofoundry analyze-study \ - Increase concurrency for throughput only before freezing a confirmatory protocol. Concurrency does not grant extra scheduled decisions, but it can change wall time and may expose backend-specific sampling behavior. + +## 8. Optional serving switches for other local servers + +Two `llm` options, both off by default, help with servers whose Responses endpoint is +incomplete or whose models copy prompt placeholders into ID fields. With both off, request +bodies, prompts and the recorded `action_schema_sha256` are exactly as before. + +```yaml +llm: + transport: chat_completions # default "responses" + extra_body: + chat_template_kwargs: {enable_thinking: false} + strict_id_patterns: true # default false +``` + +- `transport: chat_completions` posts to `{base_url}/chat/completions` with + `response_format: {type: json_schema, json_schema: {name, schema, strict}}`, the same schema + object the Responses transport sends in `text.format`. Use it when a server's `/v1/responses` + drops `text.format` (llama-server does) or does not exist. `max_tokens`, `temperature` and + `reasoning_effort` map one-to-one; usage is recorded under the Responses names so traces and + the compute audit stay uniform; `finish_reason: length` raises the same incomplete error, now + with the length, head and tail of the truncated reply; a leading `...` block is + stripped if a server returns reasoning inline. +- `extra_body` merges arbitrary fields into the request body of either transport, for + server-specific switches such as vLLM's `chat_template_kwargs`. +- `strict_id_patterns` adds a regex `pattern` to the three free-string ID fields of the action + schema (`reply_to`, `target_agent_id`, `target_artifact_id`) so constrained decoding cannot + emit placeholders such as `NONE`. It changes `action_schema_sha256`, so traces made with it + are self-labelled. The backend must honour `pattern`: on vLLM use + `--structured-outputs-config '{"backend":"guidance"}'` (xgrammar rejects the schema). diff --git a/src/biofoundry/cli.py b/src/biofoundry/cli.py index b95f114..1b738ce 100644 --- a/src/biofoundry/cli.py +++ b/src/biofoundry/cli.py @@ -195,6 +195,7 @@ def _reproducibility_metadata(config: GameConfig) -> dict[str, Any]: allow_replies=replies_available(config), resource_names=scenario.resource_names if scenario is not None else None, operation_names=scenario.operation_names if scenario is not None else None, + strict_id_patterns=config.llm.strict_id_patterns, )["schema"] prompt_text = ( SCENARIO_SYSTEM_PROMPT.format( diff --git a/src/biofoundry/config.py b/src/biofoundry/config.py index 66f940e..35753b6 100644 --- a/src/biofoundry/config.py +++ b/src/biofoundry/config.py @@ -127,6 +127,18 @@ class LLMConfig: # Work around constrained decoders that cannot terminate JSON number tokens. # Numbers travel as short decimal strings and are coerced before validation. grammar_safe_numbers: bool = False + # Replication-package addition (default off; see replication/patches/README.md). Small + # open-weight models copy the prompt's "use NONE when irrelevant" into the free-string ID + # fields, and the engine then rejects the whole action. With this flag the JSON schema + # constrains reply_to / target_agent_id / target_artifact_id to their ID grammar or "". + strict_id_patterns: bool = False + # Replication-package addition (default "responses"). "chat_completions" posts to + # {base_url}/chat/completions with response_format json_schema, for servers whose Responses + # route does not enforce text.format (llama-server converts /v1/responses internally and + # drops the schema). extra_body is merged into every request, e.g. + # {"chat_template_kwargs": {"enable_thinking": false}} for Qwen-style thinking models. + transport: str = "responses" + extra_body: dict[str, Any] | None = None # Some local tokenizers strongly prefer legal JSON whitespace forever under a # grammar. A negative bias on tokenizer tokens that decode only to whitespace # preserves compact JSON while leaving content tokens unconstrained. diff --git a/src/biofoundry/providers/openai_compatible.py b/src/biofoundry/providers/openai_compatible.py index 9d238ac..27c96ac 100644 --- a/src/biofoundry/providers/openai_compatible.py +++ b/src/biofoundry/providers/openai_compatible.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import copy import http.client import json import os @@ -94,6 +95,52 @@ def _whitespace_token_ids( return tuple(token_ids) +def _chat_usage_as_responses(usage: dict[str, Any]) -> dict[str, Any]: + """Map Chat Completions usage keys onto the Responses names the traces use.""" + mapped: dict[str, Any] = {} + for src, dst in ( + ("prompt_tokens", "input_tokens"), + ("completion_tokens", "output_tokens"), + ("total_tokens", "total_tokens"), + ): + if isinstance(usage.get(src), int): + mapped[dst] = usage[src] + return mapped + + +def _extract_chat_text(data: dict[str, Any]) -> str: + """Extract assistant text from a non-streaming Chat Completions payload.""" + error = data.get("error") + if error: + detail = str(error.get("message") or error) if isinstance(error, dict) else str(error) + raise ProviderError(f"model response failed: {detail[:1200]}") + choices = data.get("choices") + if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): + raise ProviderError(f"Chat Completions payload contains no choices: {data}") + choice = choices[0] + finish = choice.get("finish_reason") + message = choice.get("message") or {} + if finish == "length": + # Keep the head and tail of the truncated text in the error so the trace shows what + # the model was doing when it hit max_tokens (whitespace loop, repeated field, ...). + partial = message.get("content") if isinstance(message.get("content"), str) else "" + tail = "" + if partial: + tail = f" | {len(partial)} chars, head {partial[:120]!r}, tail {partial[-160:]!r}" + raise ProviderError( + "model response status is 'incomplete': {'reason': 'max_output_tokens'}" + tail + ) + if message.get("refusal"): + raise ProviderError(f"model refused the request: {str(message['refusal'])[:1200]}") + content = message.get("content") + if not isinstance(content, str) or not content: + raise ProviderError(f"Chat Completions payload contains no content: {data}") + # Thinking models without a server-side reasoning parser prepend their reasoning block. + if content.lstrip().startswith("") and "" in content: + content = content.split("", 1)[1] + return content.strip() + + def _extract_output_text(data: dict[str, Any]) -> str: """Extract assistant text from a non-streaming Responses API payload.""" error = data.get("error") @@ -159,7 +206,11 @@ def __init__( self.allow_replies = allow_replies self.resource_names = None if resource_names is None else tuple(resource_names) self.operation_names = None if operation_names is None else tuple(operation_names) - self.endpoint = config.base_url.rstrip("/") + "/responses" + if config.transport not in ("responses", "chat_completions"): + raise ValueError("llm.transport must be 'responses' or 'chat_completions'") + self.endpoint = config.base_url.rstrip("/") + ( + "/chat/completions" if config.transport == "chat_completions" else "/responses" + ) self.semaphore = asyncio.Semaphore(config.concurrency) self._json_whitespace_logit_bias: dict[str, float] = {} if config.json_whitespace_logit_bias is not None: @@ -185,6 +236,8 @@ async def generate_record( return await asyncio.to_thread(self._generate_sync_record, messages) def _request_body(self, messages: Sequence[ChatMessage]) -> dict[str, Any]: + if self.config.transport == "chat_completions": + return self._chat_completions_body(messages) body = { "model": self.config.model, "input": [ @@ -208,12 +261,58 @@ def _request_body(self, messages: Sequence[ChatMessage]) -> dict[str, Any]: allow_replies=self.allow_replies, resource_names=self.resource_names, operation_names=self.operation_names, + strict_id_patterns=self.config.strict_id_patterns, ) body["text"] = {"format": output_format} if self._json_whitespace_logit_bias: body["logit_bias"] = dict(self._json_whitespace_logit_bias) elif self.config.json_mode: body["text"] = {"format": {"type": "json_object"}} + if self.config.extra_body: + body.update(copy.deepcopy(self.config.extra_body)) + return body + + def _output_format(self) -> dict[str, Any]: + return bounded_action_plan_text_format( + self.config.max_plan_actions, + grammar_safe_numbers=self.config.grammar_safe_numbers, + allowed_action_types=self.allowed_action_types, + allow_addressing=self.allow_addressing, + allow_replies=self.allow_replies, + resource_names=self.resource_names, + operation_names=self.operation_names, + strict_id_patterns=self.config.strict_id_patterns, + ) + + def _chat_completions_body(self, messages: Sequence[ChatMessage]) -> dict[str, Any]: + """Replication-package addition: the same request as Chat Completions.""" + body: dict[str, Any] = { + "model": self.config.model, + "messages": [ + {"role": message.role, "content": message.content} for message in messages + ], + "max_tokens": self.config.max_tokens, + } + if self.config.reasoning_effort is not None: + body["reasoning_effort"] = self.config.reasoning_effort + else: + body["temperature"] = self.config.temperature + if self.config.structured_output: + output_format = self._output_format() + body["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": output_format["name"], + "schema": output_format["schema"], + "strict": bool(output_format.get("strict", True)), + }, + } + if self._json_whitespace_logit_bias: + body["logit_bias"] = dict(self._json_whitespace_logit_bias) + elif self.config.json_mode: + body["response_format"] = {"type": "json_object"} + if self.config.extra_body: + body.update(copy.deepcopy(self.config.extra_body)) return body def _generate_sync(self, messages: Sequence[ChatMessage]) -> str: @@ -262,8 +361,13 @@ def _generate_sync_record( usage = data.get("usage", {}) if not isinstance(usage, dict): usage = {} + if self.config.transport == "chat_completions": + usage = _chat_usage_as_responses(usage) + text = _extract_chat_text(data) + else: + text = _extract_output_text(data) return GenerationResult( - text=_extract_output_text(data), + text=text, usage={str(key): value for key, value in usage.items()}, metadata={ "response_id": str(data.get("id", "")), diff --git a/src/biofoundry/structured_output.py b/src/biofoundry/structured_output.py index 6c94dcf..945a36f 100644 --- a/src/biofoundry/structured_output.py +++ b/src/biofoundry/structured_output.py @@ -621,6 +621,7 @@ def bounded_action_plan_text_format( allow_replies: bool = True, resource_names: Iterable[str] | None = None, operation_names: Iterable[str] | None = None, + strict_id_patterns: bool = False, ) -> dict[str, Any]: """Return the exact provider contract for one experimental capability set.""" @@ -676,9 +677,24 @@ def bounded_action_plan_text_format( if not retained: raise ValueError("allowed_action_types must retain at least one action") result["schema"]["$defs"]["action"]["anyOf"] = retained + if strict_id_patterns: + # Replication-package addition: pin free-string ID fields to their ID grammar or the + # empty string so constrained decoders cannot emit placeholders such as "NONE". + for branch in result["schema"]["$defs"]["action"]["anyOf"]: + for field, prefix in ID_FIELD_PREFIXES.items(): + prop = branch["properties"].get(field) + if isinstance(prop, dict) and "enum" not in prop: + prop["pattern"] = f"^({prefix}_[0-9]+)?$" return result +ID_FIELD_PREFIXES = { + "reply_to": "message", + "target_agent_id": "agent", + "target_artifact_id": "artifact", +} + + def coerce_transport_numbers(value: Any) -> Any: """Convert grammar-safe numeric strings to canonical plan numbers in-place.""" diff --git a/tests/test_local_model_serving.py b/tests/test_local_model_serving.py new file mode 100644 index 0000000..da72aa8 --- /dev/null +++ b/tests/test_local_model_serving.py @@ -0,0 +1,103 @@ +"""Tests for the opt-in local-model serving options: chat-completions transport, extra_body, +constrained ID fields, and the truncated-reply detail on finish_reason=length.""" + +import hashlib +import json + +import pytest + +from biofoundry.config import LLMConfig +from biofoundry.providers.base import ChatMessage +from biofoundry.providers.openai_compatible import ( + OpenAICompatibleProvider, + ProviderError, + _chat_usage_as_responses, + _extract_chat_text, +) +from biofoundry.structured_output import ACTION_PLAN_SCHEMA, bounded_action_plan_text_format + + +def _schema_hash(**kwargs: object) -> str: + schema = bounded_action_plan_text_format(12, **kwargs)["schema"] + return hashlib.sha256(json.dumps(schema, sort_keys=True).encode()).hexdigest() + + +def test_strict_id_patterns_off_leaves_schema_unchanged() -> None: + assert _schema_hash() == _schema_hash(strict_id_patterns=False) + + +def test_strict_id_patterns_constrains_id_fields_and_changes_hash() -> None: + assert _schema_hash(strict_id_patterns=True) != _schema_hash() + schema = bounded_action_plan_text_format(12, strict_id_patterns=True)["schema"] + text = json.dumps(schema) + assert "^(message_[0-9]+)?$" in text + assert "^(agent_[0-9]+)?$" in text + assert "^(artifact_[0-9]+)?$" in text + # the plain schema carries no such patterns + assert "message_[0-9]" not in json.dumps(bounded_action_plan_text_format(12)["schema"]) + + +def test_chat_completions_transport_builds_openai_chat_body() -> None: + provider = OpenAICompatibleProvider( + LLMConfig( + model="local-qwen", + max_tokens=4096, + structured_output=True, + transport="chat_completions", + extra_body={"chat_template_kwargs": {"enable_thinking": False}}, + ) + ) + messages = [ChatMessage("system", "Act."), ChatMessage("user", "Return JSON.")] + body = provider._request_body(messages) + assert provider.endpoint == "http://127.0.0.1:8000/v1/chat/completions" + assert body["messages"] == [ + {"role": "system", "content": "Act."}, + {"role": "user", "content": "Return JSON."}, + ] + assert body["max_tokens"] == 4096 + assert body["temperature"] == 0.7 + assert body["response_format"] == { + "type": "json_schema", + "json_schema": { + "name": "biofoundry_agent_plan", + "strict": True, + "schema": ACTION_PLAN_SCHEMA, + }, + } + assert body["chat_template_kwargs"] == {"enable_thinking": False} + assert "input" not in body and "text" not in body and "store" not in body + + +def test_default_transport_is_unchanged_responses_shape() -> None: + provider = OpenAICompatibleProvider( + LLMConfig(model="gpt-5.6-luna", max_tokens=64, structured_output=True) + ) + body = provider._request_body([ChatMessage("user", "x")]) + assert provider.endpoint.endswith("/responses") + assert "input" in body and "response_format" not in body + + +def test_chat_usage_maps_to_responses_names() -> None: + usage = {"prompt_tokens": 12, "completion_tokens": 3, "total_tokens": 15} + assert _chat_usage_as_responses(usage) == { + "input_tokens": 12, + "output_tokens": 3, + "total_tokens": 15, + } + + +def test_chat_text_extraction_and_think_block_strip() -> None: + content = '\nplan\n\n{"verb":3}' + data = {"choices": [{"finish_reason": "stop", "message": {"content": content}}]} + assert _extract_chat_text(data) == '{"verb":3}' + + +def test_chat_length_finish_raises_incomplete_with_tail() -> None: + partial = "{" + "x" * 300 + data = {"choices": [{"finish_reason": "length", "message": {"content": partial}}]} + with pytest.raises(ProviderError) as info: + _extract_chat_text(data) + message = str(info.value) + assert "max_output_tokens" in message + assert "301 chars" in message + assert partial[:20] in message and partial[-20:] in message