mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-28 06:44:51 -06:00
feat: add lifecycle-owned rerank runtime
This commit is contained in:
+38
-23
@@ -846,23 +846,23 @@ boundary.
|
||||
|
||||
```
|
||||
ChatSession
|
||||
|
|
||||
+-- ResolvedModelBinding
|
||||
| +-- ModelLane (provider, client, model, capabilities, params)
|
||||
| +-- immutable ModelConfig snapshot
|
||||
| +-- registry generation
|
||||
| |
|
||||
| v
|
||||
| model_turn(ModelLane, list[Turn])
|
||||
| +-- lowering.py: Turn IR -> repaired provider-neutral wire dicts
|
||||
| +-- LLMProvider.create_streaming() (single model transport call site)
|
||||
| +--- OpenAIProvider --- OpenAI/vLLM/llama.cpp
|
||||
| +--- AnthropicProvider --- Anthropic Messages API
|
||||
| +--- GoogleProvider --- Gemini via /v1beta/openai/
|
||||
|
|
||||
v
|
||||
model_turn(ModelLane, list[Turn])
|
||||
|
|
||||
+-- lowering.py: Turn IR -> repaired provider-neutral wire dicts
|
||||
|
|
||||
v
|
||||
LLMProvider.create_streaming() (the single transport call site)
|
||||
|
|
||||
+--- OpenAIProvider --- OpenAI, vLLM, llama.cpp, any /v1/chat/completions API
|
||||
+--- AnthropicProvider --- Anthropic Messages API (native streaming, thinking)
|
||||
+--- GoogleProvider --- Google Gemini via /v1beta/openai/ (extends OpenAIProvider)
|
||||
+-- RerankLane (runtime, alias admission, registry/config witnesses)
|
||||
|
|
||||
v
|
||||
rerank(RerankLane, query, documents) --> Cohere/Jina-compatible endpoint
|
||||
```
|
||||
|
||||
**Protocol methods:**
|
||||
@@ -882,6 +882,7 @@ LLMProvider.create_streaming() (the single transport call site)
|
||||
| `StreamChunk` | `content_delta`, `reasoning_delta`, `tool_call_deltas`, `info_delta`, `usage`, `finish_reason`, `provider_blocks` |
|
||||
| `CompletionResult` | `content`, `tool_calls`, `finish_reason`, `usage`, `provider_blocks` |
|
||||
| `ModelLane` | Frozen per-loop provider/client/model binding, capabilities, sampling knobs, registry reference, and backend-auth seam |
|
||||
| `RerankLane` | Frozen per-batch binding to a registry-owned rerank runtime, stable alias admission gate, and registry/config version witnesses |
|
||||
| `ResolvedModelBinding` | A `ModelLane`, its immutable `ModelConfig`, and the registry generation read in the same snapshot |
|
||||
| `ModelTurnResult` | Canonical assistant `Turn`, tool-call dispatch mirror, serving-lane provenance, usage, and exact lowered wire facts |
|
||||
| `ModelCapabilities` | `context_window`, `max_output_tokens`, `supports_temperature`, `token_param`, `thinking_mode`, `supports_effort`, `supports_web_search`, `supports_tool_search`, `supports_vision`, `supports_reasoning_replay`, `supports_verbosity`, `verbosity`, `supports_pro_mode`, `reasoning_mode` |
|
||||
@@ -895,6 +896,17 @@ old endpoint/client state and new capabilities/configuration. Per-call operator
|
||||
toggles that are intentionally live, such as reasoning replay, are re-read by
|
||||
`model_turn()` through the lane's registry reference.
|
||||
|
||||
`RerankLane` is a deliberately narrow sibling rather than a subtype of
|
||||
`ModelLane`: endpoint reranking does not construct an LLM provider or SDK
|
||||
client. `ModelRegistry` owns the active rerank runtime, whose HTTP pool,
|
||||
per-endpoint circuit breaker, and active-call retirement state are shared by
|
||||
all sessions. Each retrieval batch resolves the Reranker alias and instruction
|
||||
from one ConfigStore snapshot and binds the matching registry generation.
|
||||
Relevant configuration changes retire the old runtime, reject new work on its
|
||||
stale lanes, and close its transport after active calls drain; cap-only reloads
|
||||
resize the stable alias admission gate without discarding the pool. Calibration
|
||||
remains an isolated one-shot client and always closes it after the probe.
|
||||
|
||||
**OpenAIProvider** (`_openai.py`): passes messages through unchanged (they are
|
||||
already in OpenAI format), including multi-part content blocks (text + images)
|
||||
in tool results. Model capability lookup covers GPT-5 through GPT-5.6,
|
||||
@@ -1001,17 +1013,20 @@ remain warm, but the session still receives a new coherent config/lane.
|
||||
simultaneous generations for that alias in one process. Every registry-backed
|
||||
role carries the same gate on its `ModelLane`, so main turns, judges, task
|
||||
agents, perception, compaction, and background generation coordinate through
|
||||
one FIFO. Two aliases never share a gate implicitly, even when their URLs are
|
||||
identical. For calls with context-first request admission, `model_turn()` admits
|
||||
the request before materializing attachment fallbacks; an oversized refusal
|
||||
therefore invokes neither attachment resolution nor nested perception/audio.
|
||||
Ordinary calls preserve their lowering cadence. Both paths finish attachment
|
||||
materialization before acquiring model capacity, then hold one lease across
|
||||
eager stream creation and the complete drain, releasing before retry backoff.
|
||||
Capacity wait is credited out of deadline accounting, preventing queued judges
|
||||
from spending their request budget before dispatch. The gate survives cap-only
|
||||
reloads in place; the field is excluded from semantic `ModelConfig` equality so
|
||||
a capacity edit does not reset judges or output-guard state.
|
||||
one FIFO. Endpoint reranking carries the same gate on its `RerankLane`, so a
|
||||
reranker alias cannot exceed that cap or bypass capacity shared with another
|
||||
role. Two aliases never share a gate implicitly, even when their URLs are
|
||||
identical. For calls with context-first request admission, `model_turn()`
|
||||
admits the request before materializing attachment fallbacks; an oversized
|
||||
refusal therefore invokes neither attachment resolution nor nested
|
||||
perception/audio. Ordinary calls preserve their lowering cadence. Both paths
|
||||
finish attachment materialization before acquiring model capacity, then hold
|
||||
one lease across eager stream creation and the complete drain, releasing before
|
||||
retry backoff. Capacity wait is credited out of deadline accounting, preventing
|
||||
queued judges from spending their request budget before dispatch. The gate
|
||||
survives cap-only reloads in place; the field is excluded from semantic
|
||||
`ModelConfig` equality so a capacity edit does not reset judges or output-guard
|
||||
state.
|
||||
|
||||
Primary loops, recursive compaction, judges, title generation, audio, and task
|
||||
agents all consume `ModelLane` rather than inspecting provider/client handles.
|
||||
|
||||
+9
-2
@@ -68,8 +68,9 @@ point to the same URL; Turnstone does not infer shared capacity from endpoint
|
||||
text. Queue time is excluded from judge/output-guard deadline accounting, and
|
||||
each retry releases its slot before backoff and reacquires for the next wire
|
||||
attempt. The cap is local to each process, not cluster-wide; account for the
|
||||
number of nodes targeting the same inference server. Direct STT/TTS protocol
|
||||
calls and Cohere/Jina reranking do not currently consume this generation cap.
|
||||
number of nodes targeting the same inference server. Cohere/Jina reranking
|
||||
selected through the Reranker role consumes the same gate as every other use of
|
||||
that alias. Direct STT/TTS protocol calls do not consume this generation cap.
|
||||
|
||||
### Judge batch parallelism
|
||||
|
||||
@@ -479,6 +480,12 @@ The alias's admission gate is retained and resized in place, so a concurrency
|
||||
edit preserves in-flight accounting and does not reset cached judges or the
|
||||
output-guard rate limiter.
|
||||
|
||||
The Reranker role is resolved from one coherent ConfigStore snapshot for each
|
||||
retrieval batch, so existing sessions observe alias or instruction changes
|
||||
without reconstruction. A relevant model-definition edit retires the old
|
||||
reranker transport and lets active requests drain on their immutable lanes;
|
||||
cap-only and unrelated edits keep the pooled transport warm.
|
||||
|
||||
Sampling and other saved workstream configuration remain workstream state. A
|
||||
model-definition edit does not silently rewrite a live workstream's chosen
|
||||
temperature, reasoning effort, max tokens, skill, or persona. Use
|
||||
|
||||
@@ -360,6 +360,16 @@ Search the web using a text query.
|
||||
|
||||
The `rerank_web_search` toggle defaults on once a reranker is selected. If the endpoint is unreachable or errors, web_search falls back silently to the backend's native result order — reranking never makes a search fail.
|
||||
|
||||
The selected alias's `max_concurrency` limit applies to reranking. Turnstone
|
||||
keeps one process-owned HTTP connection pool for the active endpoint, shares it
|
||||
across sessions, and rotates it when relevant model or Reranker-role settings
|
||||
change. Three consecutive endpoint failures open a 30-second circuit so later
|
||||
retrievals preserve native order immediately instead of repeatedly waiting for
|
||||
the network timeout; one probe is allowed after cooldown. Queued work observes
|
||||
generation cancellation. A synchronous HTTP request already dispatched cannot
|
||||
be interrupted safely, but its per-call timeout still bounds eventual drain
|
||||
and fallback latency.
|
||||
|
||||
When `rerank_bm25` is enabled, Turnstone also sends the current query and BM25 candidate metadata to the rerank endpoint. Live memory-pointer candidates contain only the memory name and authored description—never the body. Tool and skill candidates contain their names and descriptive metadata. A self-hosted endpoint (vLLM/TEI/llama.cpp) keeps this on your infrastructure; a hosted provider (Cohere/Jina/Voyage) sends it off-box.
|
||||
|
||||
**Serving a Qwen3-Reranker with vLLM.** The model is instruction-aware, so vLLM **must** apply its chat template — pass `--chat-template` explicitly. Without it the bare query produces near-random scores and reranking actively *hurts* retrieval (verified: an irrelevant passage outscored the correct one):
|
||||
|
||||
@@ -70,6 +70,36 @@ def _bootstrap_app(**overrides: Any) -> Any:
|
||||
return SimpleNamespace(state=SimpleNamespace(**state_kwargs))
|
||||
|
||||
|
||||
def test_partial_coord_teardown_shuts_registry_after_adapter() -> None:
|
||||
from turnstone.console import server as server_module
|
||||
|
||||
calls: list[str] = []
|
||||
adapter = MagicMock()
|
||||
adapter.shutdown.side_effect = lambda: calls.append("adapter")
|
||||
registry = MagicMock()
|
||||
registry.shutdown.side_effect = lambda: calls.append("registry")
|
||||
app = _bootstrap_app(coord_adapter=adapter, coord_registry=registry)
|
||||
|
||||
server_module._teardown_partial_coord_subsystem(app)
|
||||
|
||||
assert calls == ["adapter", "registry"]
|
||||
assert app.state.coord_registry is None
|
||||
|
||||
|
||||
def test_console_lifespan_shuts_coord_registry_after_adapter() -> None:
|
||||
"""Pin the normal teardown ordering without booting the full console."""
|
||||
import inspect
|
||||
|
||||
from turnstone.console import server as server_module
|
||||
|
||||
source = inspect.getsource(server_module._lifespan)
|
||||
adapter = source.find('getattr(app.state, "coord_adapter", None)')
|
||||
registry = source.find('getattr(app.state, "coord_registry", None)', adapter)
|
||||
state_writer = source.find('getattr(app.state, "coord_state_writer", None)', registry)
|
||||
assert adapter != -1 and registry != -1 and state_writer != -1
|
||||
assert adapter < registry < state_writer
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -3978,6 +3978,133 @@ class TestCancelledSendCleanupOwnership:
|
||||
session.cancel()
|
||||
raise GenerationCancelled()
|
||||
|
||||
def test_cancel_cleanup_skips_remote_rerank_and_finishes(self, tmp_db) -> None:
|
||||
"""A set cancel event cannot re-enter reranking from its own repair path."""
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.memory import save_structured_memory_strict
|
||||
from turnstone.core.rerank import RerankLane, RerankRuntime
|
||||
|
||||
class _NeverBackend:
|
||||
calls = 0
|
||||
|
||||
def rerank(self, *args: Any, **kwargs: Any):
|
||||
self.calls += 1
|
||||
raise AssertionError("cancel cleanup dispatched reranking")
|
||||
|
||||
ui = NullUI()
|
||||
session = _make_registered_session(ui=ui, user_id="owner")
|
||||
session._title_generated = True
|
||||
save_structured_memory_strict(
|
||||
"kafka_runbook",
|
||||
"private body",
|
||||
description="queued for next seam",
|
||||
scope="global",
|
||||
)
|
||||
backend = _NeverBackend()
|
||||
lane = RerankLane(
|
||||
RerankRuntime(backend, alias="rr", model="m"),
|
||||
"rr",
|
||||
"m",
|
||||
ModelAdmission("rr"),
|
||||
0,
|
||||
)
|
||||
|
||||
try:
|
||||
with (
|
||||
patch.object(session, "_resolve_rerank_lane", return_value=lane) as resolve,
|
||||
patch.object(
|
||||
session,
|
||||
"_stream_response",
|
||||
side_effect=lambda _generation: self._cancel_with_pending_cleanup(session),
|
||||
),
|
||||
):
|
||||
session.send("unrelated opening request")
|
||||
finally:
|
||||
lane.runtime.retire()
|
||||
|
||||
resolve.assert_not_called()
|
||||
assert backend.calls == 0
|
||||
assert any(turn.role is Role.TOOL for turn in session.messages)
|
||||
assert session.messages[-2].role is Role.USER
|
||||
assert session.messages[-2].text == "queued for next seam"
|
||||
assert session.messages[-1].role is Role.SYSTEM
|
||||
assert session.messages[-1].source == "memory_pointer"
|
||||
assert "kafka_runbook" in session.messages[-1].text
|
||||
assert ui.states[-1] == "idle"
|
||||
|
||||
def test_queued_turn_rerank_is_fenced_by_generation_owner(self, tmp_db) -> None:
|
||||
"""A force successor prevents the abandoned queue planner from dispatching."""
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.memory import save_structured_memory_strict
|
||||
from turnstone.core.rerank import RerankHit, RerankLane, RerankRuntime
|
||||
|
||||
class _CountingBackend:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def rerank(self, _query: str, documents: list[str], **_kwargs: Any):
|
||||
self.calls += 1
|
||||
return [RerankHit(index=i, score=1.0) for i in range(len(documents))]
|
||||
|
||||
session = _make_registered_session(user_id="owner")
|
||||
save_structured_memory_strict(
|
||||
"kafka_runbook",
|
||||
"private body",
|
||||
description="restart kafka brokers",
|
||||
scope="global",
|
||||
)
|
||||
backend = _CountingBackend()
|
||||
lane = RerankLane(
|
||||
RerankRuntime(backend, alias="rr", model="m"),
|
||||
"rr",
|
||||
"m",
|
||||
ModelAdmission("rr"),
|
||||
0,
|
||||
)
|
||||
resolve_started = threading.Event()
|
||||
release_resolve = threading.Event()
|
||||
send_errors: list[BaseException] = []
|
||||
|
||||
def resolve_lane() -> RerankLane:
|
||||
resolve_started.set()
|
||||
if not release_resolve.wait(2):
|
||||
raise RuntimeError("rerank lane resolution was not released")
|
||||
return lane
|
||||
|
||||
def stream():
|
||||
yield StreamChunk(content_delta="done")
|
||||
with session._queued_lock:
|
||||
session._queued_messages["queued-old"] = (
|
||||
"restart kafka brokers",
|
||||
"normal",
|
||||
)
|
||||
yield StreamChunk(finish_reason="stop")
|
||||
|
||||
arm_session(session, stream())
|
||||
|
||||
def run_send() -> None:
|
||||
try:
|
||||
session.send("zzzxxyy")
|
||||
except BaseException as exc:
|
||||
send_errors.append(exc)
|
||||
|
||||
sender = threading.Thread(target=run_send)
|
||||
try:
|
||||
with patch.object(session, "_resolve_rerank_lane", side_effect=resolve_lane):
|
||||
sender.start()
|
||||
assert resolve_started.wait(2)
|
||||
assert session._claim_generation() == 2
|
||||
release_resolve.set()
|
||||
sender.join(2)
|
||||
finally:
|
||||
release_resolve.set()
|
||||
sender.join(2)
|
||||
lane.runtime.retire()
|
||||
|
||||
assert not sender.is_alive()
|
||||
assert send_errors == []
|
||||
assert backend.calls == 0
|
||||
|
||||
def test_successor_waits_for_complete_cancel_cleanup_transaction(self, tmp_db):
|
||||
"""A claim already waiting on the lock observes every cleanup effect."""
|
||||
ui = NullUI()
|
||||
|
||||
@@ -1019,6 +1019,227 @@ class TestResolveEnvVars:
|
||||
assert _resolve_env_vars("") == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry-owned rerank lanes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRerankLaneRegistry:
|
||||
@staticmethod
|
||||
def _cfg(
|
||||
alias: str = "rr",
|
||||
*,
|
||||
url: str = "http://rerank.example/rerank",
|
||||
model: str = "bge",
|
||||
key: str = "secret",
|
||||
max_concurrency: int = 2,
|
||||
) -> ModelConfig:
|
||||
return ModelConfig(
|
||||
alias,
|
||||
url,
|
||||
key,
|
||||
model,
|
||||
capabilities={"supports_rerank": True},
|
||||
max_concurrency=max_concurrency,
|
||||
)
|
||||
|
||||
def test_resolve_reuses_runtime_and_stable_admission_without_llm_client(self) -> None:
|
||||
cfg = self._cfg()
|
||||
reg = ModelRegistry({"rr": cfg}, "rr")
|
||||
|
||||
first = reg.resolve_rerank_lane("rr", instruction="rank", config_version=4)
|
||||
second = reg.resolve_rerank_lane("rr", instruction="rank", config_version=5)
|
||||
|
||||
assert first.runtime is second.runtime
|
||||
assert first.admission is second.admission is reg.get_admission("rr")
|
||||
assert first.config_version == 4
|
||||
assert second.config_version == 5
|
||||
assert first.admission.limit == 2
|
||||
assert reg._clients == {}
|
||||
assert reg._providers == {}
|
||||
reg.shutdown()
|
||||
|
||||
def test_instruction_change_rotates_and_closes_old_runtime(self) -> None:
|
||||
reg = ModelRegistry({"rr": self._cfg()}, "rr")
|
||||
old = reg.resolve_rerank_lane("rr", instruction="old")
|
||||
|
||||
new = reg.resolve_rerank_lane("rr", instruction="new")
|
||||
|
||||
assert new.runtime is not old.runtime
|
||||
assert old.runtime.snapshot().retired
|
||||
assert old.runtime.snapshot().closed
|
||||
assert not new.runtime.snapshot().retired
|
||||
reg.shutdown()
|
||||
|
||||
def test_cap_only_reload_preserves_runtime_and_resizes_gate(self) -> None:
|
||||
reg = ModelRegistry({"rr": self._cfg(max_concurrency=1)}, "rr")
|
||||
old = reg.resolve_rerank_lane("rr", instruction="rank")
|
||||
|
||||
reg.reload(
|
||||
{"rr": self._cfg(max_concurrency=4)},
|
||||
"rr",
|
||||
app_state=_KEYED_STATE,
|
||||
)
|
||||
new = reg.resolve_rerank_lane("rr", instruction="rank")
|
||||
|
||||
assert new.runtime is old.runtime
|
||||
assert new.admission is old.admission
|
||||
assert new.admission.limit == 4
|
||||
assert new.registry_generation == 1
|
||||
assert not old.runtime.snapshot().retired
|
||||
reg.shutdown()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"replacement",
|
||||
[
|
||||
{"url": "http://other.example/rerank"},
|
||||
{"model": "other"},
|
||||
{"key": "different"},
|
||||
],
|
||||
)
|
||||
def test_relevant_reload_eagerly_retires_runtime(self, replacement: dict[str, str]) -> None:
|
||||
reg = ModelRegistry({"rr": self._cfg()}, "rr")
|
||||
old = reg.resolve_rerank_lane("rr")
|
||||
|
||||
reg.reload(
|
||||
{"rr": self._cfg(**replacement)},
|
||||
"rr",
|
||||
app_state=_KEYED_STATE,
|
||||
)
|
||||
|
||||
assert old.runtime.snapshot().retired
|
||||
assert old.runtime.snapshot().closed
|
||||
assert reg._rerank_runtimes == {}
|
||||
new = reg.resolve_rerank_lane("rr")
|
||||
assert new.runtime is not old.runtime
|
||||
reg.shutdown()
|
||||
|
||||
def test_relevant_reload_lets_active_call_drain_before_close(self) -> None:
|
||||
from turnstone.core.rerank import RerankHit, rerank
|
||||
|
||||
class _BlockingClient:
|
||||
def __init__(self) -> None:
|
||||
self.entered = threading.Event()
|
||||
self.release = threading.Event()
|
||||
self.close_calls = 0
|
||||
|
||||
def rerank(self, query: str, documents: list[str], **kwargs: Any) -> list[RerankHit]:
|
||||
del query, kwargs
|
||||
self.entered.set()
|
||||
assert self.release.wait(5)
|
||||
return [RerankHit(i, 1.0) for i in range(len(documents))]
|
||||
|
||||
def close(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
client = _BlockingClient()
|
||||
reg = ModelRegistry({"rr": self._cfg()}, "rr")
|
||||
with patch("turnstone.core.rerank.resolve_rerank_client", return_value=client):
|
||||
old = reg.resolve_rerank_lane("rr")
|
||||
|
||||
outcome: list[list[RerankHit]] = []
|
||||
worker = threading.Thread(
|
||||
target=lambda: outcome.append(rerank(old, "q", ["d"], timeout=2.0)),
|
||||
daemon=True,
|
||||
)
|
||||
worker.start()
|
||||
assert client.entered.wait(2)
|
||||
|
||||
reg.reload(
|
||||
{"rr": self._cfg(url="http://other.example/rerank")},
|
||||
"rr",
|
||||
app_state=_KEYED_STATE,
|
||||
)
|
||||
assert old.runtime.snapshot().retired
|
||||
assert not old.runtime.snapshot().closed
|
||||
assert client.close_calls == 0
|
||||
|
||||
client.release.set()
|
||||
worker.join(2)
|
||||
assert not worker.is_alive()
|
||||
assert outcome and outcome[0][0].index == 0
|
||||
assert old.runtime.snapshot().closed
|
||||
assert client.close_calls == 1
|
||||
reg.shutdown()
|
||||
|
||||
def test_resolving_new_role_alias_retires_previous_alias(self) -> None:
|
||||
a = self._cfg("a", url="http://a.example/rerank")
|
||||
b = self._cfg("b", url="http://b.example/rerank")
|
||||
reg = ModelRegistry({"a": a, "b": b}, "a")
|
||||
old = reg.resolve_rerank_lane("a")
|
||||
|
||||
current = reg.resolve_rerank_lane("b")
|
||||
|
||||
assert old.runtime.snapshot().closed
|
||||
assert set(reg._rerank_runtimes) == {"b"}
|
||||
assert current.admission is reg.get_admission("b")
|
||||
reg.shutdown()
|
||||
|
||||
def test_deactivate_and_shutdown_are_idempotent(self) -> None:
|
||||
reg = ModelRegistry({"rr": self._cfg()}, "rr")
|
||||
lane = reg.resolve_rerank_lane("rr")
|
||||
|
||||
reg.deactivate_rerank_runtime()
|
||||
reg.deactivate_rerank_runtime()
|
||||
reg.shutdown()
|
||||
reg.shutdown()
|
||||
|
||||
assert lane.runtime.snapshot().retired
|
||||
assert lane.runtime.snapshot().closed
|
||||
assert reg._rerank_runtimes == {}
|
||||
|
||||
@pytest.mark.parametrize("operation", ["reload", "shutdown"])
|
||||
def test_rerank_transport_closes_when_llm_client_close_raises(
|
||||
self,
|
||||
operation: str,
|
||||
) -> None:
|
||||
class _RerankClient:
|
||||
def __init__(self) -> None:
|
||||
self.close_calls = 0
|
||||
|
||||
def close(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
class _FailingLLMClient:
|
||||
def close(self) -> None:
|
||||
raise RuntimeError("llm close failed")
|
||||
|
||||
rr_cfg = self._cfg()
|
||||
llm_cfg = ModelConfig("llm", "http://old.example/v1", "key", "model")
|
||||
reg = ModelRegistry({"llm": llm_cfg, "rr": rr_cfg}, "llm")
|
||||
rerank_client = _RerankClient()
|
||||
with patch(
|
||||
"turnstone.core.rerank.resolve_rerank_client",
|
||||
return_value=rerank_client,
|
||||
):
|
||||
lane = reg.resolve_rerank_lane("rr")
|
||||
reg._clients["llm"] = _FailingLLMClient()
|
||||
|
||||
with pytest.raises(RuntimeError, match="llm close failed"):
|
||||
if operation == "reload":
|
||||
reg.reload(
|
||||
{
|
||||
"llm": dataclasses.replace(
|
||||
llm_cfg,
|
||||
base_url="http://new.example/v1",
|
||||
),
|
||||
"rr": dataclasses.replace(
|
||||
rr_cfg,
|
||||
base_url="http://new-rerank.example/rerank",
|
||||
),
|
||||
},
|
||||
"llm",
|
||||
app_state=_KEYED_STATE,
|
||||
)
|
||||
else:
|
||||
reg.shutdown()
|
||||
|
||||
assert rerank_client.close_calls == 1
|
||||
assert lane.runtime.snapshot().retired
|
||||
assert lane.runtime.snapshot().closed
|
||||
assert reg._rerank_runtimes == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ModelRegistry.reload
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+130
-54
@@ -19,17 +19,11 @@ from turnstone.core.rerank import (
|
||||
from turnstone.core.session import ChatSession
|
||||
|
||||
|
||||
def _mock_httpx_post(handler):
|
||||
"""Patch target for ``rerank.httpx.post`` that routes the call through a real
|
||||
``httpx.MockTransport``. The request flows through genuine httpx JSON/header
|
||||
encoding and response parsing — a true boundary, not a bare MagicMock.
|
||||
"""
|
||||
def _install_mock_httpx_client(monkeypatch, handler):
|
||||
"""Install the lifecycle-owned client over a real MockTransport boundary."""
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler))
|
||||
|
||||
def _post(url, **kwargs):
|
||||
return client.post(url, **kwargs)
|
||||
|
||||
return _post
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.Client", lambda: client)
|
||||
return client
|
||||
|
||||
|
||||
# A Cohere/Jina/vLLM-shaped response: results wrapper + relevance_score, returned
|
||||
@@ -112,7 +106,7 @@ class TestCohereJinaRerankClient:
|
||||
captured["auth"] = request.headers.get("authorization")
|
||||
return httpx.Response(200, json=RESULTS_WRAPPED)
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
client = CohereJinaRerankClient(
|
||||
"http://vllm:8000/rerank", model="bge", api_key="secret", timeout=10
|
||||
)
|
||||
@@ -137,7 +131,7 @@ class TestCohereJinaRerankClient:
|
||||
captured["auth"] = request.headers.get("authorization")
|
||||
return httpx.Response(200, json={"results": []})
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
CohereJinaRerankClient("http://h/rerank").rerank("q", ["a"])
|
||||
|
||||
assert captured["body"] == {"query": "q", "documents": ["a"]}
|
||||
@@ -150,7 +144,7 @@ class TestCohereJinaRerankClient:
|
||||
called["n"] += 1
|
||||
return httpx.Response(200, json={"results": []})
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
assert CohereJinaRerankClient("http://h/rerank").rerank("q", []) == []
|
||||
assert called["n"] == 0 # short-circuits before any HTTP call
|
||||
|
||||
@@ -158,7 +152,7 @@ class TestCohereJinaRerankClient:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=BARE_LIST)
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
hits = CohereJinaRerankClient("http://tei/rerank").rerank("q", ["a", "b"])
|
||||
assert [h.index for h in hits] == [1, 0]
|
||||
|
||||
@@ -166,7 +160,7 @@ class TestCohereJinaRerankClient:
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(401, text="unauthorized")
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
CohereJinaRerankClient("http://h/rerank").rerank("q", ["a"])
|
||||
|
||||
@@ -208,6 +202,7 @@ class _FakeConfigStore:
|
||||
def __init__(self, values: dict, stored: set[str] | None = None) -> None:
|
||||
self._values = dict(values)
|
||||
self._stored = frozenset(stored if stored is not None else values.keys())
|
||||
self.version = 1
|
||||
|
||||
def stored_keys(self) -> frozenset[str]:
|
||||
return self._stored
|
||||
@@ -215,11 +210,14 @@ class _FakeConfigStore:
|
||||
def get(self, key: str):
|
||||
return self._values.get(key)
|
||||
|
||||
def effective_snapshot(self):
|
||||
return self.version, dict(self._values)
|
||||
|
||||
|
||||
class TestSessionRerankWiring:
|
||||
def test_disabled_when_no_config_store(self):
|
||||
stub = SimpleNamespace(_config_store=None, _registry=None, tool_timeout=30)
|
||||
assert ChatSession._resolve_rerank_client(stub) is None
|
||||
assert ChatSession._resolve_rerank_lane(stub) is None
|
||||
|
||||
def test_disabled_when_no_reranker_alias(self):
|
||||
# No Reranker role selected -> reranking off. There is no global endpoint
|
||||
@@ -230,7 +228,7 @@ class TestSessionRerankWiring:
|
||||
tool_timeout=30,
|
||||
_registry=SimpleNamespace(get_config=lambda a: None),
|
||||
)
|
||||
assert ChatSession._resolve_rerank_client(stub) is None
|
||||
assert ChatSession._resolve_rerank_lane(stub) is None
|
||||
|
||||
def test_enabled_for_defaults_true_without_store(self):
|
||||
stub = SimpleNamespace(_config_store=None)
|
||||
@@ -254,16 +252,18 @@ class TestSessionRerankWiring:
|
||||
capabilities={"supports_rerank": True},
|
||||
)
|
||||
cs = _FakeConfigStore({"tools.reranker_alias": "rr"})
|
||||
stub = SimpleNamespace(
|
||||
_config_store=cs,
|
||||
tool_timeout=30,
|
||||
_registry=SimpleNamespace(get_config=lambda a: cfg),
|
||||
)
|
||||
client = ChatSession._resolve_rerank_client(stub)
|
||||
from turnstone.core.model_registry import ModelRegistry
|
||||
|
||||
registry = ModelRegistry({"rr": cfg}, "rr")
|
||||
stub = SimpleNamespace(_config_store=cs, tool_timeout=30, _registry=registry)
|
||||
lane = ChatSession._resolve_rerank_lane(stub)
|
||||
assert lane is not None
|
||||
client = lane.runtime.client
|
||||
assert isinstance(client, CohereJinaRerankClient)
|
||||
assert client._url == "http://rr:8000/rerank"
|
||||
assert client._model == "bge"
|
||||
assert client._api_key == "k"
|
||||
registry.shutdown()
|
||||
|
||||
def test_ignores_alias_without_rerank_capability(self):
|
||||
# A non-reranker model (no supports_rerank) must NOT be used as a reranker,
|
||||
@@ -274,12 +274,58 @@ class TestSessionRerankWiring:
|
||||
alias="chat", base_url="http://chat/v1", api_key="k", model="gpt", capabilities={}
|
||||
)
|
||||
cs = _FakeConfigStore({"tools.reranker_alias": "chat"})
|
||||
stub = SimpleNamespace(
|
||||
_config_store=cs,
|
||||
tool_timeout=30,
|
||||
_registry=SimpleNamespace(get_config=lambda a: cfg),
|
||||
from turnstone.core.model_registry import ModelRegistry
|
||||
|
||||
registry = ModelRegistry({"chat": cfg}, "chat")
|
||||
stub = SimpleNamespace(_config_store=cs, tool_timeout=30, _registry=registry)
|
||||
assert ChatSession._resolve_rerank_lane(stub) is None
|
||||
registry.shutdown()
|
||||
|
||||
def test_unrelated_settings_version_reuses_runtime(self):
|
||||
from turnstone.core.model_registry import ModelConfig, ModelRegistry
|
||||
|
||||
cfg = ModelConfig(
|
||||
"rr",
|
||||
"http://rr/rerank",
|
||||
"k",
|
||||
"bge",
|
||||
capabilities={"supports_rerank": True},
|
||||
)
|
||||
assert ChatSession._resolve_rerank_client(stub) is None
|
||||
cs = _FakeConfigStore({"tools.reranker_alias": "rr", "tools.rerank_instruction": "rank"})
|
||||
registry = ModelRegistry({"rr": cfg}, "rr")
|
||||
stub = SimpleNamespace(_config_store=cs, _registry=registry)
|
||||
|
||||
first = ChatSession._resolve_rerank_lane(stub)
|
||||
cs._values["unrelated.setting"] = True
|
||||
cs.version += 1
|
||||
second = ChatSession._resolve_rerank_lane(stub)
|
||||
|
||||
assert first is not None and second is not None
|
||||
assert first.runtime is second.runtime
|
||||
assert first.config_version != second.config_version
|
||||
registry.shutdown()
|
||||
|
||||
def test_clearing_role_retires_previous_runtime_on_next_resolution(self):
|
||||
from turnstone.core.model_registry import ModelConfig, ModelRegistry
|
||||
|
||||
cfg = ModelConfig(
|
||||
"rr",
|
||||
"http://rr/rerank",
|
||||
"k",
|
||||
"bge",
|
||||
capabilities={"supports_rerank": True},
|
||||
)
|
||||
cs = _FakeConfigStore({"tools.reranker_alias": "rr"})
|
||||
registry = ModelRegistry({"rr": cfg}, "rr")
|
||||
stub = SimpleNamespace(_config_store=cs, _registry=registry)
|
||||
lane = ChatSession._resolve_rerank_lane(stub)
|
||||
assert lane is not None
|
||||
|
||||
cs._values["tools.reranker_alias"] = ""
|
||||
cs.version += 1
|
||||
assert ChatSession._resolve_rerank_lane(stub) is None
|
||||
assert lane.runtime.snapshot().closed
|
||||
registry.shutdown()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -287,18 +333,6 @@ class TestSessionRerankWiring:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeRerankClient:
|
||||
"""In-process RerankClient stub returning fixed hits (no HTTP)."""
|
||||
|
||||
def __init__(self, hits: list[RerankHit]) -> None:
|
||||
self._hits = hits
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], *, top_n: int | None = None
|
||||
) -> list[RerankHit]:
|
||||
return self._hits
|
||||
|
||||
|
||||
class TestRerankInstruction:
|
||||
"""`rerank_instruction` wraps the query for instruction-aware rerankers.
|
||||
|
||||
@@ -318,13 +352,13 @@ class TestRerankInstruction:
|
||||
|
||||
def test_no_instruction_sends_bare_query(self, monkeypatch):
|
||||
sent, handler = self._capture()
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
CohereJinaRerankClient("http://x/rerank").rerank("capital of France", ["d0"])
|
||||
assert sent["body"]["query"] == "capital of France"
|
||||
|
||||
def test_instruction_wraps_query(self, monkeypatch):
|
||||
sent, handler = self._capture()
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
_install_mock_httpx_client(monkeypatch, handler)
|
||||
CohereJinaRerankClient("http://x/rerank", instruction="Find relevant passages").rerank(
|
||||
"capital of France", ["d0"]
|
||||
)
|
||||
@@ -374,31 +408,35 @@ class TestNormalizeScores:
|
||||
class TestSessionBM25Reranker:
|
||||
"""``_bm25_reranker`` / ``_bm25_rerank_threshold`` — the BM25 seam adapters.
|
||||
|
||||
Drives the real closure through a fake ``RerankClient`` (the in-process
|
||||
callable seam), never by patching internal state. The HTTP boundary lives in
|
||||
``TestCohereJinaRerankClient`` above.
|
||||
Drives the real closure through fixed hits at the session dispatch seam.
|
||||
The HTTP boundary lives in ``TestCohereJinaRerankClient`` above.
|
||||
"""
|
||||
|
||||
def test_none_when_disabled_for_bm25(self):
|
||||
# Per-tool toggle off -> no reranker even with an endpoint configured.
|
||||
def test_disabled_bm25_is_identity_order(self):
|
||||
# Long-lived closures stay installed across settings changes; disabled
|
||||
# state is an identity rerank and must not resolve a runtime.
|
||||
stub = SimpleNamespace(
|
||||
_rerank_enabled_for=lambda tool: False,
|
||||
_resolve_rerank_client=lambda: _FakeRerankClient([]),
|
||||
_rerank_hits=lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("disabled reranker dispatched")
|
||||
),
|
||||
)
|
||||
assert ChatSession._bm25_reranker(stub) is None
|
||||
rank = ChatSession._bm25_reranker(stub)
|
||||
assert rank("q", ["d0", "d1"]) == [0, 1]
|
||||
|
||||
def test_none_when_no_client(self):
|
||||
# Enabled, but no endpoint resolves -> None.
|
||||
def test_no_lane_is_identity_order(self):
|
||||
# Enabled, but no endpoint resolves: preserve BM25 pool order.
|
||||
stub = SimpleNamespace(
|
||||
_rerank_enabled_for=lambda tool: True,
|
||||
_resolve_rerank_client=lambda: None,
|
||||
_rerank_hits=lambda *args, **kwargs: None,
|
||||
)
|
||||
assert ChatSession._bm25_reranker(stub) is None
|
||||
rank = ChatSession._bm25_reranker(stub)
|
||||
assert rank("q", ["d0", "d1"]) == [0, 1]
|
||||
|
||||
def _enabled_stub(self, hits: list[RerankHit]) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
_rerank_enabled_for=lambda tool: True,
|
||||
_resolve_rerank_client=lambda: _FakeRerankClient(hits),
|
||||
_rerank_hits=lambda query, docs, origin_generation=0: hits,
|
||||
)
|
||||
|
||||
def test_no_threshold_returns_all_hit_indices(self):
|
||||
@@ -461,6 +499,44 @@ class TestSessionBM25Reranker:
|
||||
assert rank is not None
|
||||
assert rank("q", ["d0", "d1"]) == [0]
|
||||
|
||||
def test_generation_cancellation_bypasses_bm25_exception_fallback(self):
|
||||
import threading
|
||||
from types import MethodType
|
||||
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.bm25 import BM25Index
|
||||
from turnstone.core.rerank import RerankLane, RerankRuntime
|
||||
from turnstone.core.session import GenerationCancelled
|
||||
|
||||
class _NeverBackend:
|
||||
def rerank(self, *args, **kwargs):
|
||||
raise AssertionError("superseded generation dispatched")
|
||||
|
||||
lane = RerankLane(
|
||||
RerankRuntime(_NeverBackend(), alias="rr", model="m"),
|
||||
"rr",
|
||||
"m",
|
||||
ModelAdmission("rr"),
|
||||
0,
|
||||
)
|
||||
stub = SimpleNamespace(
|
||||
_generation=2,
|
||||
_generation_lock=threading.Lock(),
|
||||
_cancel_event=threading.Event(),
|
||||
_publication_shutdown=False,
|
||||
_rerank_enabled_for=lambda tool: True,
|
||||
_resolve_rerank_lane=lambda: lane,
|
||||
tool_timeout=10,
|
||||
)
|
||||
stub._check_cancelled = MethodType(ChatSession._check_cancelled, stub)
|
||||
stub._rerank_hits = MethodType(ChatSession._rerank_hits, stub)
|
||||
rank = ChatSession._bm25_reranker(stub, origin_generation=1)
|
||||
|
||||
with pytest.raises(GenerationCancelled):
|
||||
BM25Index(["alpha"], reranker=rank).search("alpha")
|
||||
assert lane.runtime.snapshot().circuit.consecutive_failures == 0
|
||||
lane.runtime.retire()
|
||||
|
||||
def test_threshold_reads_setting(self):
|
||||
cs = _FakeConfigStore({"tools.rerank_bm25_threshold": 0.42})
|
||||
stub = SimpleNamespace(_config_store=cs)
|
||||
|
||||
@@ -2,8 +2,16 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.rerank import RerankHit
|
||||
from turnstone.core.rerank_calibrate import _GAP_FRACTION, _PROBE_SET, _build_result, calibrate
|
||||
from turnstone.core.rerank_calibrate import (
|
||||
_GAP_FRACTION,
|
||||
_PROBE_SET,
|
||||
_build_result,
|
||||
calibrate,
|
||||
calibrate_model,
|
||||
)
|
||||
|
||||
|
||||
class _ScriptedClient:
|
||||
@@ -95,6 +103,42 @@ class TestCalibrate:
|
||||
assert res.raw_scale == "unknown (no scores)"
|
||||
|
||||
|
||||
class TestCalibrateModelLifecycle:
|
||||
class _ClosableClient(_ScriptedClient):
|
||||
def __init__(self, *, fail: bool = False) -> None:
|
||||
super().__init__(0.9, 0.1)
|
||||
self.fail = fail
|
||||
self.close_calls = 0
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], *, top_n: int | None = None
|
||||
) -> list[RerankHit]:
|
||||
if self.fail:
|
||||
raise RuntimeError("probe failed")
|
||||
return super().rerank(query, documents, top_n=top_n)
|
||||
|
||||
def close(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
def test_one_shot_client_closes_after_success(self, monkeypatch) -> None:
|
||||
client = self._ClosableClient()
|
||||
monkeypatch.setattr("turnstone.core.rerank.resolve_rerank_client", lambda *a, **k: client)
|
||||
|
||||
result = calibrate_model("http://rr/rerank", "m", "k")
|
||||
|
||||
assert result.separated
|
||||
assert client.close_calls == 1
|
||||
|
||||
def test_one_shot_client_closes_after_failure(self, monkeypatch) -> None:
|
||||
client = self._ClosableClient(fail=True)
|
||||
monkeypatch.setattr("turnstone.core.rerank.resolve_rerank_client", lambda *a, **k: client)
|
||||
|
||||
with pytest.raises(RuntimeError, match="probe failed"):
|
||||
calibrate_model("http://rr/rerank", "m", "k")
|
||||
|
||||
assert client.close_calls == 1
|
||||
|
||||
|
||||
class TestCalibrationCapsConfinement:
|
||||
def test_calibrate_merge_confined_to_calibration_fields(self):
|
||||
"""The merge touches only the three probe-derived keys and preserves
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
"""Live lifecycle and admission checks for the managed rerank runtime.
|
||||
|
||||
The local counting proxy preserves real endpoint responses while making the
|
||||
client-side pool and admission boundary observable. Run explicitly against a
|
||||
Cohere/Jina-compatible endpoint::
|
||||
|
||||
TURNSTONE_LIVE_RERANK_URL=http://127.0.0.1:8000/rerank \
|
||||
TURNSTONE_LIVE_RERANK_MODEL=my-reranker \
|
||||
pytest tests/test_rerank_live.py -m live -v
|
||||
|
||||
``TURNSTONE_LIVE_RERANK_API_KEY`` is optional.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from turnstone.core.model_registry import ModelConfig, ModelRegistry
|
||||
from turnstone.core.rerank import RerankHit, RerankLane, rerank
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from types import TracebackType
|
||||
|
||||
|
||||
_LIVE_URL = os.environ.get("TURNSTONE_LIVE_RERANK_URL", "").strip()
|
||||
_LIVE_MODEL = os.environ.get("TURNSTONE_LIVE_RERANK_MODEL", "").strip()
|
||||
_LIVE_API_KEY = os.environ.get("TURNSTONE_LIVE_RERANK_API_KEY", "").strip()
|
||||
_HOP_BY_HOP_HEADERS = frozenset(
|
||||
{
|
||||
"connection",
|
||||
"content-encoding",
|
||||
"content-length",
|
||||
"host",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
"te",
|
||||
"trailer",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Observed:
|
||||
requests: int
|
||||
active: int
|
||||
peak: int
|
||||
accepted_connections: int
|
||||
|
||||
|
||||
class _Counter:
|
||||
"""Count proxy requests and optionally hold the first admitted cohort."""
|
||||
|
||||
def __init__(self, *, rendezvous_size: int = 1, hold_seconds: float = 0.0) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._rendezvous = threading.Event()
|
||||
self._rendezvous_size = rendezvous_size
|
||||
self._hold_seconds = hold_seconds
|
||||
self._requests = 0
|
||||
self._active = 0
|
||||
self._peak = 0
|
||||
self._accepted_connections = 0
|
||||
if rendezvous_size <= 1:
|
||||
self._rendezvous.set()
|
||||
|
||||
def accepted(self) -> None:
|
||||
with self._lock:
|
||||
self._accepted_connections += 1
|
||||
|
||||
def enter(self) -> None:
|
||||
with self._lock:
|
||||
self._requests += 1
|
||||
self._active += 1
|
||||
self._peak = max(self._peak, self._active)
|
||||
if self._active >= self._rendezvous_size:
|
||||
self._rendezvous.set()
|
||||
assert self._rendezvous.wait(10), "live rerank requests did not overlap"
|
||||
if self._hold_seconds:
|
||||
time.sleep(self._hold_seconds)
|
||||
|
||||
def leave(self) -> None:
|
||||
with self._lock:
|
||||
self._active -= 1
|
||||
|
||||
def snapshot(self) -> _Observed:
|
||||
with self._lock:
|
||||
return _Observed(
|
||||
requests=self._requests,
|
||||
active=self._active,
|
||||
peak=self._peak,
|
||||
accepted_connections=self._accepted_connections,
|
||||
)
|
||||
|
||||
|
||||
class _LiveProxyServer(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self, counter: _Counter) -> None:
|
||||
self.counter = counter
|
||||
headers = {"Authorization": f"Bearer {_LIVE_API_KEY}"} if _LIVE_API_KEY else {}
|
||||
self.upstream = httpx.Client(
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(60.0, connect=5.0),
|
||||
)
|
||||
super().__init__(("127.0.0.1", 0), _LiveProxyHandler)
|
||||
|
||||
def get_request(self):
|
||||
request, address = super().get_request()
|
||||
self.counter.accepted()
|
||||
return request, address
|
||||
|
||||
|
||||
class _LiveProxyHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
@property
|
||||
def _proxy(self) -> _LiveProxyServer:
|
||||
if not isinstance(self.server, _LiveProxyServer):
|
||||
raise TypeError("live rerank handler requires _LiveProxyServer")
|
||||
return self.server
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802 - stdlib handler contract
|
||||
length = int(self.headers.get("content-length", "0"))
|
||||
body = self.rfile.read(length)
|
||||
self._proxy.counter.enter()
|
||||
try:
|
||||
headers = {
|
||||
name: value
|
||||
for name, value in self.headers.items()
|
||||
if name.lower() not in _HOP_BY_HOP_HEADERS and name.lower() != "authorization"
|
||||
}
|
||||
response = self._proxy.upstream.post(_LIVE_URL, content=body, headers=headers)
|
||||
payload = response.content
|
||||
self.send_response(response.status_code)
|
||||
content_type = response.headers.get("content-type")
|
||||
if content_type:
|
||||
self.send_header("Content-Type", content_type)
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
self.wfile.flush()
|
||||
finally:
|
||||
self._proxy.counter.leave()
|
||||
|
||||
def log_message(self, _format: str, *args: Any) -> None:
|
||||
del args
|
||||
|
||||
|
||||
class _LiveProxy:
|
||||
def __init__(self, counter: _Counter) -> None:
|
||||
self._server = _LiveProxyServer(counter)
|
||||
self._thread = threading.Thread(
|
||||
target=self._server.serve_forever,
|
||||
name="rerank-live-proxy",
|
||||
daemon=True,
|
||||
)
|
||||
|
||||
@property
|
||||
def url(self) -> str:
|
||||
host, port = self._server.server_address
|
||||
return f"http://{host}:{port}/rerank"
|
||||
|
||||
def __enter__(self) -> _LiveProxy:
|
||||
self._thread.start()
|
||||
return self
|
||||
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
del exc_type, exc_value, traceback
|
||||
self._server.shutdown()
|
||||
self._server.server_close()
|
||||
self._server.upstream.close()
|
||||
self._thread.join(timeout=5)
|
||||
if self._thread.is_alive():
|
||||
raise AssertionError("live rerank proxy did not stop")
|
||||
|
||||
|
||||
def _registry_lane(url: str, *, limit: int) -> tuple[ModelRegistry, RerankLane]:
|
||||
cfg = ModelConfig(
|
||||
alias="live-reranker",
|
||||
base_url=url,
|
||||
api_key="",
|
||||
model=_LIVE_MODEL,
|
||||
max_concurrency=limit,
|
||||
capabilities={"supports_rerank": True},
|
||||
)
|
||||
registry = ModelRegistry({cfg.alias: cfg}, default=cfg.alias)
|
||||
return registry, registry.resolve_rerank_lane(cfg.alias)
|
||||
|
||||
|
||||
def _dispatch(lane: RerankLane) -> list[RerankHit]:
|
||||
return rerank(
|
||||
lane,
|
||||
"What is the capital of France?",
|
||||
["Paris is the capital of France.", "Whales are mammals."],
|
||||
timeout=30.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.live
|
||||
@pytest.mark.skipif(
|
||||
not _LIVE_URL or not _LIVE_MODEL,
|
||||
reason="TURNSTONE_LIVE_RERANK_URL and TURNSTONE_LIVE_RERANK_MODEL are required",
|
||||
)
|
||||
class TestLiveRerankRuntime:
|
||||
def test_real_scores_and_reuses_one_client_connection(self) -> None:
|
||||
counter = _Counter()
|
||||
with _LiveProxy(counter) as proxy:
|
||||
registry, first_lane = _registry_lane(proxy.url, limit=2)
|
||||
try:
|
||||
second_lane = registry.resolve_rerank_lane("live-reranker")
|
||||
assert second_lane.runtime is first_lane.runtime
|
||||
first = _dispatch(first_lane)
|
||||
second = _dispatch(second_lane)
|
||||
finally:
|
||||
registry.shutdown()
|
||||
|
||||
assert first[0].index == second[0].index == 0
|
||||
assert all(math.isfinite(hit.score) for hit in [*first, *second])
|
||||
observed = counter.snapshot()
|
||||
assert observed.requests == 2
|
||||
assert observed.active == 0
|
||||
assert observed.accepted_connections == 1
|
||||
|
||||
def test_real_endpoint_never_exceeds_alias_cap(self) -> None:
|
||||
counter = _Counter(rendezvous_size=2, hold_seconds=0.25)
|
||||
with _LiveProxy(counter) as proxy:
|
||||
registry, lane = _registry_lane(proxy.url, limit=2)
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=4) as pool:
|
||||
futures = [pool.submit(_dispatch, lane) for _ in range(4)]
|
||||
results = [future.result(timeout=75) for future in futures]
|
||||
finally:
|
||||
registry.shutdown()
|
||||
|
||||
assert all(result and result[0].index == 0 for result in results)
|
||||
observed = counter.snapshot()
|
||||
assert observed.requests == 4
|
||||
assert observed.active == 0
|
||||
assert observed.peak == 2
|
||||
@@ -0,0 +1,408 @@
|
||||
"""Lifecycle, admission, circuit, and real keep-alive tests for RerankLane."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.deadline import DeadlineCancelledError, StreamAbortRef
|
||||
from turnstone.core.rerank import (
|
||||
CohereJinaRerankClient,
|
||||
RerankCircuitOpenError,
|
||||
RerankHit,
|
||||
RerankLane,
|
||||
RerankRuntime,
|
||||
RerankRuntimeRetiredError,
|
||||
rerank,
|
||||
)
|
||||
|
||||
|
||||
class _Backend:
|
||||
"""Thread-safe scripted backend with close and concurrency accounting."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
self.active = 0
|
||||
self.max_active = 0
|
||||
self.close_calls = 0
|
||||
self.fail = False
|
||||
self.entered = threading.Event()
|
||||
self.release: threading.Event | None = None
|
||||
self.on_call: Any = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def rerank(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
*,
|
||||
top_n: int | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> list[RerankHit]:
|
||||
del query, top_n, timeout
|
||||
with self._lock:
|
||||
self.calls += 1
|
||||
self.active += 1
|
||||
self.max_active = max(self.max_active, self.active)
|
||||
self.entered.set()
|
||||
try:
|
||||
if self.on_call is not None:
|
||||
self.on_call()
|
||||
if self.release is not None:
|
||||
assert self.release.wait(5), "test backend release timed out"
|
||||
if self.fail:
|
||||
raise RuntimeError("rerank endpoint failed")
|
||||
return [RerankHit(index=i, score=1.0 - i / 100) for i in range(len(documents))]
|
||||
finally:
|
||||
with self._lock:
|
||||
self.active -= 1
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
self.close_calls += 1
|
||||
|
||||
|
||||
def _lane(
|
||||
backend: _Backend,
|
||||
*,
|
||||
limit: int = 0,
|
||||
failure_threshold: int = 3,
|
||||
cooldown_seconds: float = 30.0,
|
||||
clock: Any = time.monotonic,
|
||||
) -> RerankLane:
|
||||
runtime = RerankRuntime(
|
||||
backend,
|
||||
alias="rr",
|
||||
model="m",
|
||||
failure_threshold=failure_threshold,
|
||||
cooldown_seconds=cooldown_seconds,
|
||||
clock=clock,
|
||||
)
|
||||
return RerankLane(runtime, "rr", "m", ModelAdmission("rr", limit), 0)
|
||||
|
||||
|
||||
def _call(lane: RerankLane, cancel_ref: Any = None) -> list[RerankHit]:
|
||||
return rerank(lane, "q", ["d0", "d1"], timeout=2.0, cancel_ref=cancel_ref)
|
||||
|
||||
|
||||
class TestRuntimeLifecycle:
|
||||
def test_empty_batch_touches_no_runtime_state(self) -> None:
|
||||
backend = _Backend()
|
||||
lane = _lane(backend)
|
||||
|
||||
assert rerank(lane, "q", [], timeout=1.0) == []
|
||||
snapshot = lane.runtime.snapshot()
|
||||
assert backend.calls == 0
|
||||
assert snapshot.active_calls == 0
|
||||
assert snapshot.circuit.state == "closed"
|
||||
assert snapshot.circuit.consecutive_failures == 0
|
||||
|
||||
def test_retire_idle_closes_exactly_once(self) -> None:
|
||||
backend = _Backend()
|
||||
runtime = _lane(backend).runtime
|
||||
|
||||
runtime.retire()
|
||||
runtime.retire()
|
||||
|
||||
assert backend.close_calls == 1
|
||||
assert runtime.snapshot().retired
|
||||
assert runtime.snapshot().closed
|
||||
with pytest.raises(RerankRuntimeRetiredError):
|
||||
runtime.acquire_call()
|
||||
|
||||
def test_active_call_drains_before_close_and_new_call_is_refused(self) -> None:
|
||||
backend = _Backend()
|
||||
backend.release = threading.Event()
|
||||
lane = _lane(backend)
|
||||
outcome: list[Any] = []
|
||||
|
||||
worker = threading.Thread(target=lambda: outcome.append(_call(lane)), daemon=True)
|
||||
worker.start()
|
||||
assert backend.entered.wait(2)
|
||||
|
||||
lane.runtime.retire()
|
||||
during = lane.runtime.snapshot()
|
||||
assert during.retired
|
||||
assert not during.closed
|
||||
assert backend.close_calls == 0
|
||||
with pytest.raises(RerankRuntimeRetiredError):
|
||||
_call(lane)
|
||||
assert backend.calls == 1
|
||||
|
||||
backend.release.set()
|
||||
worker.join(2)
|
||||
assert not worker.is_alive()
|
||||
assert len(outcome) == 1
|
||||
after = lane.runtime.snapshot()
|
||||
assert after.closed
|
||||
assert backend.close_calls == 1
|
||||
|
||||
|
||||
class TestAdmission:
|
||||
def test_shared_alias_cap_bounds_peak_dispatch(self) -> None:
|
||||
backend = _Backend()
|
||||
backend.release = threading.Event()
|
||||
lane = _lane(backend, limit=2)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=4) as pool:
|
||||
futures = [pool.submit(_call, lane) for _ in range(4)]
|
||||
deadline = time.monotonic() + 2
|
||||
while backend.active < 2 and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert backend.active == 2
|
||||
assert lane.admission.snapshot().queued == 2
|
||||
backend.release.set()
|
||||
assert all(len(f.result(timeout=2)) == 2 for f in futures)
|
||||
|
||||
assert backend.calls == 4
|
||||
assert backend.max_active == 2
|
||||
assert lane.admission.snapshot().in_flight == 0
|
||||
|
||||
def test_cancelled_waiter_is_removed_without_dispatch_or_circuit_failure(self) -> None:
|
||||
backend = _Backend()
|
||||
lane = _lane(backend, limit=1)
|
||||
held = lane.admission.acquire()
|
||||
cancel_ref = StreamAbortRef()
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def _waiter() -> None:
|
||||
try:
|
||||
_call(lane, cancel_ref)
|
||||
except BaseException as exc: # capture the exact cancellation type
|
||||
errors.append(exc)
|
||||
|
||||
worker = threading.Thread(target=_waiter, daemon=True)
|
||||
worker.start()
|
||||
deadline = time.monotonic() + 2
|
||||
while lane.admission.snapshot().queued != 1 and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert lane.admission.snapshot().queued == 1
|
||||
cancel_ref.abort()
|
||||
worker.join(2)
|
||||
held.release()
|
||||
|
||||
assert not worker.is_alive()
|
||||
assert len(errors) == 1
|
||||
assert isinstance(errors[0], DeadlineCancelledError)
|
||||
assert backend.calls == 0
|
||||
snapshot = lane.runtime.snapshot()
|
||||
assert snapshot.circuit.state == "closed"
|
||||
assert snapshot.circuit.consecutive_failures == 0
|
||||
|
||||
def test_post_request_cancellation_is_not_endpoint_failure(self) -> None:
|
||||
cancel_ref = StreamAbortRef()
|
||||
backend = _Backend()
|
||||
backend.on_call = cancel_ref.abort
|
||||
lane = _lane(backend)
|
||||
|
||||
with pytest.raises(DeadlineCancelledError):
|
||||
_call(lane, cancel_ref)
|
||||
|
||||
assert backend.calls == 1
|
||||
snapshot = lane.runtime.snapshot()
|
||||
assert snapshot.circuit.state == "closed"
|
||||
assert snapshot.circuit.consecutive_failures == 0
|
||||
|
||||
def test_cancellation_wins_when_dispatched_request_also_fails(self) -> None:
|
||||
cancel_ref = StreamAbortRef()
|
||||
backend = _Backend()
|
||||
backend.on_call = cancel_ref.abort
|
||||
backend.fail = True
|
||||
lane = _lane(backend)
|
||||
|
||||
with pytest.raises(DeadlineCancelledError):
|
||||
_call(lane, cancel_ref)
|
||||
|
||||
assert backend.calls == 1
|
||||
snapshot = lane.runtime.snapshot()
|
||||
assert snapshot.circuit.state == "closed"
|
||||
assert snapshot.circuit.consecutive_failures == 0
|
||||
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 100.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
class TestCircuit:
|
||||
def test_open_fast_fallback_and_successful_half_open_probe(self) -> None:
|
||||
clock = _Clock()
|
||||
backend = _Backend()
|
||||
backend.fail = True
|
||||
lane = _lane(backend, failure_threshold=3, cooldown_seconds=30, clock=clock)
|
||||
|
||||
for _ in range(3):
|
||||
with pytest.raises(RuntimeError, match="endpoint failed"):
|
||||
_call(lane)
|
||||
assert backend.calls == 3
|
||||
assert lane.runtime.snapshot().circuit.state == "open"
|
||||
|
||||
with pytest.raises(RerankCircuitOpenError):
|
||||
_call(lane)
|
||||
assert backend.calls == 3 # no HTTP/backend dispatch while open
|
||||
|
||||
clock.now += 30
|
||||
backend.fail = False
|
||||
assert len(_call(lane)) == 2
|
||||
assert backend.calls == 4
|
||||
recovered = lane.runtime.snapshot().circuit
|
||||
assert recovered.state == "closed"
|
||||
assert recovered.consecutive_failures == 0
|
||||
|
||||
def test_open_circuit_preserves_bm25_order_without_another_dispatch(self) -> None:
|
||||
from turnstone.core.bm25 import BM25Index
|
||||
|
||||
backend = _Backend()
|
||||
backend.fail = True
|
||||
lane = _lane(backend, failure_threshold=1)
|
||||
documents = ["alpha alpha", "alpha beta", "beta"]
|
||||
expected = BM25Index(documents).search("alpha", k=3)
|
||||
|
||||
def _rank(query: str, candidates: list[str]) -> list[int]:
|
||||
return [hit.index for hit in rerank(lane, query, candidates, timeout=2.0)]
|
||||
|
||||
index = BM25Index(documents, reranker=_rank)
|
||||
assert index.search("alpha", k=3) == expected
|
||||
assert backend.calls == 1
|
||||
assert index.search("alpha", k=3) == expected
|
||||
assert backend.calls == 1
|
||||
|
||||
def test_only_one_half_open_probe_dispatches(self) -> None:
|
||||
clock = _Clock()
|
||||
backend = _Backend()
|
||||
backend.fail = True
|
||||
lane = _lane(backend, failure_threshold=1, cooldown_seconds=5, clock=clock)
|
||||
with pytest.raises(RuntimeError):
|
||||
_call(lane)
|
||||
|
||||
clock.now += 5
|
||||
backend.fail = False
|
||||
backend.entered.clear()
|
||||
backend.release = threading.Event()
|
||||
outcome: list[Any] = []
|
||||
probe = threading.Thread(target=lambda: outcome.append(_call(lane)), daemon=True)
|
||||
probe.start()
|
||||
assert backend.entered.wait(2)
|
||||
|
||||
with pytest.raises(RerankCircuitOpenError):
|
||||
_call(lane)
|
||||
assert backend.calls == 2
|
||||
|
||||
backend.release.set()
|
||||
probe.join(2)
|
||||
assert not probe.is_alive()
|
||||
assert len(outcome) == 1
|
||||
assert lane.runtime.snapshot().circuit.state == "closed"
|
||||
|
||||
def test_failed_half_open_probe_starts_fresh_cooldown(self) -> None:
|
||||
clock = _Clock()
|
||||
backend = _Backend()
|
||||
backend.fail = True
|
||||
lane = _lane(backend, failure_threshold=1, cooldown_seconds=10, clock=clock)
|
||||
with pytest.raises(RuntimeError):
|
||||
_call(lane)
|
||||
|
||||
clock.now += 10
|
||||
with pytest.raises(RuntimeError):
|
||||
_call(lane)
|
||||
assert backend.calls == 2
|
||||
|
||||
clock.now += 9
|
||||
with pytest.raises(RerankCircuitOpenError):
|
||||
_call(lane)
|
||||
assert backend.calls == 2
|
||||
|
||||
def test_waiters_admitted_after_open_do_not_dispatch(self) -> None:
|
||||
backend = _Backend()
|
||||
backend.fail = True
|
||||
backend.release = threading.Event()
|
||||
lane = _lane(backend, limit=1, failure_threshold=3)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=4) as pool:
|
||||
futures = [pool.submit(_call, lane) for _ in range(4)]
|
||||
assert backend.entered.wait(2)
|
||||
deadline = time.monotonic() + 2
|
||||
while lane.admission.snapshot().queued != 3 and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert lane.admission.snapshot().queued == 3
|
||||
backend.release.set()
|
||||
errors = [future.exception(timeout=2) for future in futures]
|
||||
|
||||
assert sum(type(error) is RuntimeError for error in errors) == 3
|
||||
assert sum(isinstance(error, RerankCircuitOpenError) for error in errors) == 1
|
||||
assert backend.calls == 3
|
||||
assert lane.admission.snapshot().in_flight == 0
|
||||
assert lane.runtime.snapshot().circuit.state == "open"
|
||||
|
||||
|
||||
class _ConnectionCountingServer(ThreadingHTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.accepted_connections = 0
|
||||
self.requests = 0
|
||||
super().__init__(("127.0.0.1", 0), _KeepAliveHandler)
|
||||
|
||||
def get_request(self):
|
||||
request, address = super().get_request()
|
||||
self.accepted_connections += 1
|
||||
return request, address
|
||||
|
||||
|
||||
class _KeepAliveHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802 - stdlib handler contract
|
||||
length = int(self.headers.get("content-length", "0"))
|
||||
json.loads(self.rfile.read(length))
|
||||
server = self.server
|
||||
assert isinstance(server, _ConnectionCountingServer)
|
||||
server.requests += 1
|
||||
body = json.dumps(
|
||||
{"results": [{"index": 0, "relevance_score": 0.9}]},
|
||||
separators=(",", ":"),
|
||||
).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
self.wfile.flush()
|
||||
|
||||
def log_message(self, fmt: str, *args: Any) -> None:
|
||||
del fmt, args
|
||||
|
||||
|
||||
def test_repeated_reranks_reuse_one_real_tcp_connection() -> None:
|
||||
server = _ConnectionCountingServer()
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
url = f"http://127.0.0.1:{server.server_port}/rerank"
|
||||
client = CohereJinaRerankClient(url, model="m")
|
||||
lane = RerankLane(
|
||||
RerankRuntime(client, alias="rr", model="m"),
|
||||
"rr",
|
||||
"m",
|
||||
ModelAdmission("rr"),
|
||||
0,
|
||||
)
|
||||
try:
|
||||
assert [hit.index for hit in _call(lane)] == [0]
|
||||
assert [hit.index for hit in _call(lane)] == [0]
|
||||
assert server.requests == 2
|
||||
assert server.accepted_connections == 1
|
||||
finally:
|
||||
lane.runtime.retire()
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(2)
|
||||
@@ -106,7 +106,7 @@ def _make_session(*, kind: str = "interactive", user_id: str = "test-user") -> A
|
||||
# Truncation budget — required by _truncate_output on every exec.
|
||||
session.tool_truncation = 100_000
|
||||
# skills(action='find') ranks via BM25Index(..., reranker=self._bm25_reranker()),
|
||||
# which reaches _resolve_rerank_client -> self.tool_timeout. No _config_store/
|
||||
# which reaches _resolve_rerank_lane -> self.tool_timeout. No _config_store/
|
||||
# _registry here -> no endpoint -> reranker is None -> pure-BM25 path.
|
||||
session.tool_timeout = 30
|
||||
|
||||
|
||||
@@ -5382,6 +5382,10 @@ def _bootstrap_coord_subsystem(
|
||||
coord_adapter.shutdown()
|
||||
except Exception:
|
||||
log.warning("console.coord_bootstrap_rollback_adapter_failed", exc_info=True)
|
||||
try:
|
||||
coord_registry.shutdown()
|
||||
except Exception:
|
||||
log.warning("console.coord_bootstrap_rollback_registry_failed", exc_info=True)
|
||||
try:
|
||||
shutdown_idle_nudge_watchers(app)
|
||||
except Exception:
|
||||
@@ -5470,6 +5474,7 @@ def _load_and_bootstrap_coord_subsystem(app: Starlette, storage: Any, config_sto
|
||||
"""
|
||||
from turnstone.core.model_registry import load_model_registry
|
||||
|
||||
coord_registry: Any | None = None
|
||||
try:
|
||||
try:
|
||||
coord_registry = load_model_registry(storage=storage)
|
||||
@@ -5489,6 +5494,11 @@ def _load_and_bootstrap_coord_subsystem(app: Starlette, storage: Any, config_sto
|
||||
_bootstrap_coord_subsystem(app, storage, config_store, coord_registry)
|
||||
except Exception:
|
||||
log.warning("console.coordinator_init_failed", exc_info=True)
|
||||
if coord_registry is not None:
|
||||
try:
|
||||
coord_registry.shutdown()
|
||||
except Exception:
|
||||
log.warning("console.coord_startup_registry_shutdown_failed", exc_info=True)
|
||||
# ``_bootstrap_coord_subsystem`` rolls back its own partial
|
||||
# side-effects from locals before re-raising, so this helper
|
||||
# is normally redundant — kept as defence-in-depth in case
|
||||
@@ -5556,6 +5566,13 @@ def _teardown_partial_coord_subsystem(app: Any) -> None:
|
||||
except Exception:
|
||||
log.warning("console.coord_partial_adapter_shutdown_failed", exc_info=True)
|
||||
|
||||
registry = getattr(state, "coord_registry", None)
|
||||
if registry is not None:
|
||||
try:
|
||||
registry.shutdown()
|
||||
except Exception:
|
||||
log.warning("console.coord_partial_registry_shutdown_failed", exc_info=True)
|
||||
|
||||
# Idle nudge watchers are tracked in a list on app.state; the
|
||||
# console only ever installs one (the coord watcher), so a blanket
|
||||
# shutdown is safe — there is no other watcher to tear down by
|
||||
@@ -5875,6 +5892,15 @@ async def _lifespan(app: Starlette) -> AsyncGenerator[None, None]:
|
||||
coord_adapter_shutdown.shutdown()
|
||||
except Exception:
|
||||
log.debug("console.coord_adapter_shutdown_failed", exc_info=True)
|
||||
coord_registry_shutdown = getattr(app.state, "coord_registry", None)
|
||||
if coord_registry_shutdown is not None:
|
||||
try:
|
||||
# Coordinator sessions stop through the adapter first; the registry
|
||||
# can then retire pooled model and rerank transports without a new
|
||||
# session resolving behind the teardown.
|
||||
await asyncio.to_thread(coord_registry_shutdown.shutdown)
|
||||
except Exception:
|
||||
log.debug("console.coord_registry_shutdown_failed", exc_info=True)
|
||||
coord_state_writer_shutdown = getattr(app.state, "coord_state_writer", None)
|
||||
if coord_state_writer_shutdown is not None:
|
||||
try:
|
||||
@@ -12852,6 +12878,10 @@ def _maybe_bootstrap_coord_subsystem(app: Any, storage: Any) -> None:
|
||||
_bootstrap_coord_subsystem(app, storage, config_store, coord_registry)
|
||||
except Exception as exc:
|
||||
log.warning("console.coord_bootstrap_failed", exc_info=True)
|
||||
try:
|
||||
coord_registry.shutdown()
|
||||
except Exception:
|
||||
log.warning("console.coord_bootstrap_registry_shutdown_failed", exc_info=True)
|
||||
# Tear down any partially-stamped handles so a later retry
|
||||
# via the same CRUD path doesn't spawn duplicate daemons.
|
||||
_teardown_partial_coord_subsystem(app)
|
||||
|
||||
@@ -7,6 +7,8 @@ resilience when the primary model is unreachable.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import re
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
@@ -14,7 +16,9 @@ from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
|
||||
from turnstone.core.rerank import RerankLane, RerankRuntime
|
||||
|
||||
from turnstone.core.admission import ModelAdmission
|
||||
from turnstone.core.config import load_config
|
||||
@@ -23,6 +27,20 @@ from turnstone.core.providers import LLMProvider, create_client, create_provider
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _deferred_close_actions(actions: list[Callable[[], None]]) -> Iterator[None]:
|
||||
"""Run collected transport closes after the enclosing registry lock exits."""
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for close in actions:
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
log.warning("rerank.deferred_close_failed", exc_info=True)
|
||||
|
||||
|
||||
MODEL_AUTH_MODES = frozenset({"static", "entra_obo", "entra_app", "rfc8693_obo"})
|
||||
|
||||
# One bound for the backend-auth text columns (obo_audience, obo_scopes),
|
||||
@@ -448,6 +466,37 @@ def _validate_registry_args(
|
||||
raise ValueError(f"Task model '{task_model}' not found in registry")
|
||||
|
||||
|
||||
def _rerank_runtime_fingerprint(cfg: ModelConfig, instruction: str) -> bytes:
|
||||
"""Return a non-secret digest of configuration that changes rerank behavior."""
|
||||
caps = cfg.capabilities
|
||||
values = (
|
||||
cfg.base_url.strip(),
|
||||
cfg.api_key,
|
||||
cfg.model.strip(),
|
||||
cfg.provider,
|
||||
cfg.auth_mode,
|
||||
cfg.obo_audience,
|
||||
cfg.obo_scopes,
|
||||
str(caps.get("rerank_mode") or "endpoint"),
|
||||
str(bool(caps.get("supports_rerank"))),
|
||||
str(bool(caps.get("supports_prefill_rerank"))),
|
||||
instruction.strip(),
|
||||
)
|
||||
digest = hashlib.sha256()
|
||||
for value in values:
|
||||
encoded = value.encode("utf-8")
|
||||
digest.update(len(encoded).to_bytes(8, "big"))
|
||||
digest.update(encoded)
|
||||
return digest.digest()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _RerankRuntimeEntry:
|
||||
runtime: RerankRuntime
|
||||
fingerprint: bytes = field(repr=False)
|
||||
instruction: str = field(repr=False)
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Holds named model configurations with thread-safe lazy client creation.
|
||||
|
||||
@@ -482,6 +531,7 @@ class ModelRegistry:
|
||||
self.task_effort = task_effort
|
||||
self._clients: dict[str, Any] = {}
|
||||
self._providers: dict[str, LLMProvider] = {}
|
||||
self._rerank_runtimes: dict[str, _RerankRuntimeEntry] = {}
|
||||
self._admissions = {
|
||||
alias: ModelAdmission(alias, int(getattr(cfg, "max_concurrency", 0)))
|
||||
for alias, cfg in self._models.items()
|
||||
@@ -577,6 +627,88 @@ class ModelRegistry:
|
||||
raise UnknownModelAliasError(alias)
|
||||
return gate
|
||||
|
||||
def resolve_rerank_lane(
|
||||
self,
|
||||
alias: str,
|
||||
*,
|
||||
instruction: str = "",
|
||||
config_version: int = 0,
|
||||
) -> RerankLane:
|
||||
"""Resolve one coherent endpoint-backed rerank binding.
|
||||
|
||||
Unlike :meth:`resolve_binding`, this path deliberately constructs no
|
||||
LLM provider or SDK client: a Cohere/Jina rerank route is not a chat
|
||||
model surface. It shares only the alias's stable admission object and
|
||||
registry generation with normal model lanes.
|
||||
"""
|
||||
from turnstone.core.rerank import RerankLane, RerankRuntime, resolve_rerank_client
|
||||
|
||||
close_actions: list[Callable[[], None]] = []
|
||||
instruction = instruction.strip()
|
||||
with _deferred_close_actions(close_actions), self._client_lock:
|
||||
cfg = self._models.get(alias)
|
||||
if cfg is None:
|
||||
raise UnknownModelAliasError(alias)
|
||||
if not cfg.base_url.strip() or not cfg.capabilities.get("supports_rerank"):
|
||||
raise ValueError(f"Model alias {alias!r} is not a configured reranker")
|
||||
|
||||
fingerprint = _rerank_runtime_fingerprint(cfg, instruction)
|
||||
entry = self._rerank_runtimes.get(alias)
|
||||
if entry is None or entry.fingerprint != fingerprint:
|
||||
client = resolve_rerank_client(
|
||||
cfg.base_url,
|
||||
model=cfg.model,
|
||||
api_key=cfg.api_key,
|
||||
instruction=instruction,
|
||||
)
|
||||
if client is None: # guarded above; retain a fail-closed boundary
|
||||
raise ValueError(f"Model alias {alias!r} has no rerank endpoint")
|
||||
replacement = _RerankRuntimeEntry(
|
||||
runtime=RerankRuntime(client, alias=alias, model=cfg.model),
|
||||
fingerprint=fingerprint,
|
||||
instruction=instruction,
|
||||
)
|
||||
if entry is not None:
|
||||
close = entry.runtime.begin_retirement()
|
||||
if close is not None:
|
||||
close_actions.append(close)
|
||||
log.info("rerank.runtime_retired alias=%s reason=config", alias)
|
||||
self._rerank_runtimes[alias] = replacement
|
||||
entry = replacement
|
||||
|
||||
# The Reranker role selects one alias per process. Retire a
|
||||
# previously selected alias when a settings change resolves its
|
||||
# replacement; active calls drain on their old immutable lanes.
|
||||
for other_alias, other in list(self._rerank_runtimes.items()):
|
||||
if other_alias == alias:
|
||||
continue
|
||||
close = other.runtime.begin_retirement()
|
||||
if close is not None:
|
||||
close_actions.append(close)
|
||||
del self._rerank_runtimes[other_alias]
|
||||
log.info("rerank.runtime_retired alias=%s reason=role_change", other_alias)
|
||||
|
||||
lane = RerankLane(
|
||||
runtime=entry.runtime,
|
||||
alias=alias,
|
||||
model=cfg.model,
|
||||
admission=self._admissions[alias],
|
||||
registry_generation=self._generation,
|
||||
config_version=config_version,
|
||||
)
|
||||
return lane
|
||||
|
||||
def deactivate_rerank_runtime(self) -> None:
|
||||
"""Retire any selected rerank runtime after the role becomes empty/invalid."""
|
||||
close_actions: list[Callable[[], None]] = []
|
||||
with _deferred_close_actions(close_actions), self._client_lock:
|
||||
for alias, entry in list(self._rerank_runtimes.items()):
|
||||
close = entry.runtime.begin_retirement()
|
||||
if close is not None:
|
||||
close_actions.append(close)
|
||||
log.info("rerank.runtime_retired alias=%s reason=disabled", alias)
|
||||
self._rerank_runtimes.clear()
|
||||
|
||||
def has_alias(self, alias: str) -> bool:
|
||||
"""Check if *alias* exists in the registry."""
|
||||
return alias in self._models
|
||||
@@ -728,7 +860,8 @@ class ModelRegistry:
|
||||
# refuses on this deployment — say so at every swap chokepoint.
|
||||
warn_profile_mismatched_aliases(models, app_state)
|
||||
_validate_registry_args(models, default, fallback, agent_model, task_model)
|
||||
with self._client_lock:
|
||||
rerank_close_actions: list[Callable[[], None]] = []
|
||||
with _deferred_close_actions(rerank_close_actions), self._client_lock:
|
||||
# FIRST write inside the lock, deliberately BEFORE the map swap
|
||||
# and the client teardown. The per-send refresh reads the maps
|
||||
# lock-free and samples the generation AFTER them (see
|
||||
@@ -761,6 +894,23 @@ class ModelRegistry:
|
||||
self._admissions[alias] = ModelAdmission(alias, limit)
|
||||
else:
|
||||
gate.set_limit(limit)
|
||||
# Rerank runtimes have their own transport and breaker state. A
|
||||
# cap-only or unrelated model reload preserves them; any relevant
|
||||
# endpoint/model/auth/capability change retires the old runtime now.
|
||||
for alias, entry in list(self._rerank_runtimes.items()):
|
||||
rerank_cfg = self._models.get(alias)
|
||||
if (
|
||||
rerank_cfg is None
|
||||
or not rerank_cfg.base_url.strip()
|
||||
or not rerank_cfg.capabilities.get("supports_rerank")
|
||||
or entry.fingerprint
|
||||
!= _rerank_runtime_fingerprint(rerank_cfg, entry.instruction)
|
||||
):
|
||||
close = entry.runtime.begin_retirement()
|
||||
if close is not None:
|
||||
rerank_close_actions.append(close)
|
||||
del self._rerank_runtimes[alias]
|
||||
log.info("rerank.runtime_retired alias=%s reason=registry_reload", alias)
|
||||
# Removed aliases remain as tombstones for this registry's
|
||||
# lifetime. A stale lane may still hold or queue on that object;
|
||||
# re-adding the alias must reconfigure the same gate rather than
|
||||
@@ -800,7 +950,13 @@ class ModelRegistry:
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""Close all cached client connections."""
|
||||
with self._client_lock:
|
||||
rerank_close_actions: list[Callable[[], None]] = []
|
||||
with _deferred_close_actions(rerank_close_actions), self._client_lock:
|
||||
for entry in self._rerank_runtimes.values():
|
||||
close = entry.runtime.begin_retirement()
|
||||
if close is not None:
|
||||
rerank_close_actions.append(close)
|
||||
self._rerank_runtimes.clear()
|
||||
for client in self._clients.values():
|
||||
if hasattr(client, "close"):
|
||||
client.close()
|
||||
|
||||
+413
-4
@@ -21,15 +21,22 @@ default and no fall back to a local model.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import math
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
from turnstone.core.deadline import DeadlineCancelledError
|
||||
from turnstone.core.log import get_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from turnstone.core.admission import AdmissionLease, ModelAdmission
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
# A reranker reorders candidate documents by relevance to the query, returning
|
||||
@@ -49,6 +56,14 @@ class RerankError(RuntimeError):
|
||||
"""
|
||||
|
||||
|
||||
class RerankCircuitOpenError(RerankError):
|
||||
"""The endpoint circuit is open, so no rerank request was dispatched."""
|
||||
|
||||
|
||||
class RerankRuntimeRetiredError(RerankError):
|
||||
"""The resolved runtime retired before this batch could dispatch."""
|
||||
|
||||
|
||||
def _sigmoid(x: float) -> float:
|
||||
"""Numerically-stable logistic sigmoid (maps a logit to a probability)."""
|
||||
if x >= 0.0:
|
||||
@@ -89,12 +104,385 @@ class RerankClient(Protocol):
|
||||
"""Minimal interface for a rerank backend."""
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], *, top_n: int | None = None
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
*,
|
||||
top_n: int | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> list[RerankHit]:
|
||||
"""Score ``documents`` against ``query``; return hits sorted best-first."""
|
||||
...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RerankCircuitSnapshot:
|
||||
"""Non-sensitive instantaneous breaker state for diagnostics and tests."""
|
||||
|
||||
state: str
|
||||
consecutive_failures: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RerankRuntimeSnapshot:
|
||||
"""Non-sensitive instantaneous lifecycle state for diagnostics and tests."""
|
||||
|
||||
alias: str
|
||||
model: str
|
||||
active_calls: int
|
||||
retired: bool
|
||||
closed: bool
|
||||
circuit: RerankCircuitSnapshot
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CircuitPermit:
|
||||
epoch: int
|
||||
probe: bool
|
||||
|
||||
|
||||
class _RerankCircuit:
|
||||
"""Small thread-safe consecutive-failure circuit for one rerank runtime."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
alias: str,
|
||||
model: str,
|
||||
*,
|
||||
failure_threshold: int,
|
||||
cooldown_seconds: float,
|
||||
clock: Callable[[], float],
|
||||
) -> None:
|
||||
self._alias = alias
|
||||
self._model = model
|
||||
self._failure_threshold = failure_threshold
|
||||
self._cooldown_seconds = cooldown_seconds
|
||||
self._clock = clock
|
||||
self._lock = threading.Lock()
|
||||
self._state = "closed"
|
||||
self._consecutive_failures = 0
|
||||
self._opened_at: float | None = None
|
||||
# Invalidates completions from calls admitted before a transition.
|
||||
self._epoch = 0
|
||||
|
||||
def acquire(self) -> _CircuitPermit:
|
||||
with self._lock:
|
||||
if self._state == "closed":
|
||||
return _CircuitPermit(self._epoch, False)
|
||||
if self._state == "half_open":
|
||||
raise RerankCircuitOpenError(
|
||||
f"rerank circuit is half-open for alias {self._alias!r}"
|
||||
)
|
||||
opened_at = self._opened_at
|
||||
if opened_at is None or self._clock() - opened_at < self._cooldown_seconds:
|
||||
raise RerankCircuitOpenError(f"rerank circuit is open for alias {self._alias!r}")
|
||||
self._state = "half_open"
|
||||
log.info(
|
||||
"rerank.circuit_half_open alias=%s model=%s",
|
||||
self._alias,
|
||||
self._model,
|
||||
)
|
||||
return _CircuitPermit(self._epoch, True)
|
||||
|
||||
def validate(self, permit: _CircuitPermit) -> None:
|
||||
"""Refuse a permit invalidated while its caller waited for admission."""
|
||||
with self._lock:
|
||||
expected_state = "half_open" if permit.probe else "closed"
|
||||
if permit.epoch != self._epoch or self._state != expected_state:
|
||||
raise RerankCircuitOpenError(
|
||||
f"rerank circuit changed while waiting for alias {self._alias!r}"
|
||||
)
|
||||
|
||||
def succeed(self, permit: _CircuitPermit) -> None:
|
||||
recovered = False
|
||||
with self._lock:
|
||||
if permit.epoch != self._epoch:
|
||||
return
|
||||
if permit.probe:
|
||||
if self._state != "half_open":
|
||||
return
|
||||
self._state = "closed"
|
||||
self._opened_at = None
|
||||
self._consecutive_failures = 0
|
||||
self._epoch += 1
|
||||
recovered = True
|
||||
elif self._state == "closed":
|
||||
self._consecutive_failures = 0
|
||||
if recovered:
|
||||
log.info(
|
||||
"rerank.circuit_recovered alias=%s model=%s",
|
||||
self._alias,
|
||||
self._model,
|
||||
)
|
||||
|
||||
def fail(self, permit: _CircuitPermit) -> None:
|
||||
opened = False
|
||||
failures = 0
|
||||
with self._lock:
|
||||
if permit.epoch != self._epoch:
|
||||
return
|
||||
if permit.probe:
|
||||
if self._state != "half_open":
|
||||
return
|
||||
self._state = "open"
|
||||
self._opened_at = self._clock()
|
||||
self._consecutive_failures = self._failure_threshold
|
||||
self._epoch += 1
|
||||
failures = self._consecutive_failures
|
||||
opened = True
|
||||
elif self._state == "closed":
|
||||
self._consecutive_failures += 1
|
||||
failures = self._consecutive_failures
|
||||
if failures >= self._failure_threshold:
|
||||
self._state = "open"
|
||||
self._opened_at = self._clock()
|
||||
self._epoch += 1
|
||||
opened = True
|
||||
if opened:
|
||||
log.warning(
|
||||
"rerank.circuit_open alias=%s model=%s failures=%d",
|
||||
self._alias,
|
||||
self._model,
|
||||
failures,
|
||||
)
|
||||
|
||||
def abandon(self, permit: _CircuitPermit) -> None:
|
||||
"""Release a half-open reservation without judging endpoint health."""
|
||||
with self._lock:
|
||||
if permit.epoch == self._epoch and permit.probe and self._state == "half_open":
|
||||
# Preserve the original opened_at. Its cooldown has already
|
||||
# elapsed, so the next caller may become the replacement probe.
|
||||
self._state = "open"
|
||||
|
||||
def snapshot(self) -> RerankCircuitSnapshot:
|
||||
with self._lock:
|
||||
return RerankCircuitSnapshot(self._state, self._consecutive_failures)
|
||||
|
||||
|
||||
class _RerankRuntimeLease:
|
||||
"""One idempotently releasable active-call hold on a runtime."""
|
||||
|
||||
__slots__ = ("_released", "_runtime")
|
||||
|
||||
def __init__(self, runtime: RerankRuntime) -> None:
|
||||
self._runtime = runtime
|
||||
self._released = False
|
||||
|
||||
def release(self) -> None:
|
||||
if self._released:
|
||||
return
|
||||
self._released = True
|
||||
self._runtime._release_call()
|
||||
|
||||
|
||||
class RerankRuntime:
|
||||
"""Lifecycle-owned backend, circuit, and active-call retirement state."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: RerankClient,
|
||||
*,
|
||||
alias: str,
|
||||
model: str,
|
||||
failure_threshold: int = 3,
|
||||
cooldown_seconds: float = 30.0,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
if failure_threshold < 1:
|
||||
raise ValueError("rerank circuit failure threshold must be positive")
|
||||
if cooldown_seconds < 0:
|
||||
raise ValueError("rerank circuit cooldown must be non-negative")
|
||||
self._client = client
|
||||
self.alias = alias
|
||||
self.model = model
|
||||
self._state_lock = threading.Lock()
|
||||
self._active_calls = 0
|
||||
self._retired = False
|
||||
self._closed = False
|
||||
self._circuit = _RerankCircuit(
|
||||
alias,
|
||||
model,
|
||||
failure_threshold=failure_threshold,
|
||||
cooldown_seconds=cooldown_seconds,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
@property
|
||||
def client(self) -> RerankClient:
|
||||
return self._client
|
||||
|
||||
def acquire_call(self) -> _RerankRuntimeLease:
|
||||
with self._state_lock:
|
||||
if self._retired or self._closed:
|
||||
raise RerankRuntimeRetiredError(f"rerank runtime retired for alias {self.alias!r}")
|
||||
self._active_calls += 1
|
||||
return _RerankRuntimeLease(self)
|
||||
|
||||
def _release_call(self) -> None:
|
||||
close_now = False
|
||||
with self._state_lock:
|
||||
if self._active_calls <= 0:
|
||||
raise RuntimeError("rerank runtime lease released without an active call")
|
||||
self._active_calls -= 1
|
||||
if self._retired and self._active_calls == 0 and not self._closed:
|
||||
self._closed = True
|
||||
close_now = True
|
||||
if close_now:
|
||||
self._close_client()
|
||||
|
||||
def begin_retirement(self) -> Callable[[], None] | None:
|
||||
"""Prevent new calls and return an idle close action for the caller.
|
||||
|
||||
The split lets ``ModelRegistry`` mark the runtime while holding its
|
||||
short-lived registry lock, then execute a potentially blocking client
|
||||
close after releasing that lock. Active calls close from their final
|
||||
lease release instead.
|
||||
"""
|
||||
with self._state_lock:
|
||||
self._retired = True
|
||||
if self._active_calls or self._closed:
|
||||
return None
|
||||
self._closed = True
|
||||
return self._close_client
|
||||
|
||||
def retire(self) -> None:
|
||||
close = self.begin_retirement()
|
||||
if close is not None:
|
||||
close()
|
||||
|
||||
def _close_client(self) -> None:
|
||||
close = getattr(self._client, "close", None)
|
||||
if not callable(close):
|
||||
return
|
||||
try:
|
||||
close()
|
||||
except Exception:
|
||||
log.warning(
|
||||
"rerank.runtime_close_failed alias=%s model=%s",
|
||||
self.alias,
|
||||
self.model,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def circuit_acquire(self) -> _CircuitPermit:
|
||||
return self._circuit.acquire()
|
||||
|
||||
def circuit_succeed(self, permit: _CircuitPermit) -> None:
|
||||
self._circuit.succeed(permit)
|
||||
|
||||
def circuit_validate(self, permit: _CircuitPermit) -> None:
|
||||
self._circuit.validate(permit)
|
||||
|
||||
def circuit_fail(self, permit: _CircuitPermit) -> None:
|
||||
self._circuit.fail(permit)
|
||||
|
||||
def circuit_abandon(self, permit: _CircuitPermit) -> None:
|
||||
self._circuit.abandon(permit)
|
||||
|
||||
def snapshot(self) -> RerankRuntimeSnapshot:
|
||||
with self._state_lock:
|
||||
active_calls = self._active_calls
|
||||
retired = self._retired
|
||||
closed = self._closed
|
||||
return RerankRuntimeSnapshot(
|
||||
alias=self.alias,
|
||||
model=self.model,
|
||||
active_calls=active_calls,
|
||||
retired=retired,
|
||||
closed=closed,
|
||||
circuit=self._circuit.snapshot(),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RerankLane:
|
||||
"""One immutable, per-batch binding to a shared rerank runtime."""
|
||||
|
||||
runtime: RerankRuntime
|
||||
alias: str
|
||||
model: str
|
||||
admission: ModelAdmission
|
||||
registry_generation: int
|
||||
config_version: int = 0
|
||||
|
||||
|
||||
def _raise_if_aborted(cancel_ref: Any) -> None:
|
||||
if bool(getattr(cancel_ref, "aborted", False)):
|
||||
raise DeadlineCancelledError("rerank cancelled")
|
||||
|
||||
|
||||
def rerank(
|
||||
lane: RerankLane,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
*,
|
||||
top_n: int | None = None,
|
||||
timeout: float,
|
||||
cancel_ref: Any = None,
|
||||
) -> list[RerankHit]:
|
||||
"""Dispatch one rerank batch through circuit, admission, and runtime leases."""
|
||||
if not documents:
|
||||
return []
|
||||
_raise_if_aborted(cancel_ref)
|
||||
permit = lane.runtime.circuit_acquire()
|
||||
admission: AdmissionLease | None = None
|
||||
runtime_lease: _RerankRuntimeLease | None = None
|
||||
dispatched = False
|
||||
try:
|
||||
admission = lane.admission.acquire(cancel_ref=cancel_ref)
|
||||
# The circuit can open while this call is queued behind either rerank
|
||||
# or model work on the shared alias gate. Do not let an old closed-state
|
||||
# permit leak one more full-timeout request after that transition.
|
||||
lane.runtime.circuit_validate(permit)
|
||||
runtime_lease = lane.runtime.acquire_call()
|
||||
_raise_if_aborted(cancel_ref)
|
||||
mark_dispatch = getattr(cancel_ref, "mark_dispatch", None)
|
||||
if callable(mark_dispatch):
|
||||
with contextlib.suppress(Exception):
|
||||
mark_dispatch()
|
||||
dispatched = True
|
||||
hits = lane.runtime.client.rerank(
|
||||
query,
|
||||
documents,
|
||||
top_n=top_n,
|
||||
timeout=timeout,
|
||||
)
|
||||
_raise_if_aborted(cancel_ref)
|
||||
if not hits:
|
||||
raise RerankError("rerank endpoint returned no scores for non-empty input")
|
||||
except (DeadlineCancelledError, RerankCircuitOpenError, RerankRuntimeRetiredError):
|
||||
lane.runtime.circuit_abandon(permit)
|
||||
raise
|
||||
except Exception:
|
||||
try:
|
||||
# Stop/supersession owns the outcome even when the dispatched
|
||||
# transport fails while cancellation is landing. Do not charge a
|
||||
# cancelled request to endpoint health or let retrieval fallback
|
||||
# absorb it as an ordinary rerank outage.
|
||||
_raise_if_aborted(cancel_ref)
|
||||
except DeadlineCancelledError:
|
||||
lane.runtime.circuit_abandon(permit)
|
||||
raise
|
||||
if dispatched:
|
||||
lane.runtime.circuit_fail(permit)
|
||||
else:
|
||||
# Admission and lifecycle failures say nothing about endpoint
|
||||
# health; only a call that reached the backend can trip its circuit.
|
||||
lane.runtime.circuit_abandon(permit)
|
||||
raise
|
||||
except BaseException:
|
||||
lane.runtime.circuit_abandon(permit)
|
||||
raise
|
||||
else:
|
||||
lane.runtime.circuit_succeed(permit)
|
||||
return hits
|
||||
finally:
|
||||
if runtime_lease is not None:
|
||||
runtime_lease.release()
|
||||
if admission is not None:
|
||||
admission.release()
|
||||
|
||||
|
||||
class CohereJinaRerankClient:
|
||||
"""Rerank via a Cohere/Jina-compatible ``POST <url>`` endpoint.
|
||||
|
||||
@@ -117,9 +505,17 @@ class CohereJinaRerankClient:
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
self._instruction = instruction
|
||||
self._client = httpx.Client()
|
||||
self._close_lock = threading.Lock()
|
||||
self._closed = False
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], *, top_n: int | None = None
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
*,
|
||||
top_n: int | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> list[RerankHit]:
|
||||
if not documents:
|
||||
return []
|
||||
@@ -135,10 +531,23 @@ class CohereJinaRerankClient:
|
||||
if top_n is not None:
|
||||
payload["top_n"] = top_n
|
||||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
resp = httpx.post(self._url, json=payload, headers=headers, timeout=self._timeout)
|
||||
resp = self._client.post(
|
||||
self._url,
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=self._timeout if timeout is None else timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return _parse_hits(resp.json(), len(documents))
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the owned HTTP connection pool exactly once."""
|
||||
with self._close_lock:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
self._client.close()
|
||||
|
||||
|
||||
def _parse_hits(data: Any, n_docs: int) -> list[RerankHit]:
|
||||
"""Parse a Cohere/Jina ``{"results": [...]}`` or bare-list rerank response.
|
||||
|
||||
@@ -298,7 +298,12 @@ def calibrate_model(
|
||||
)
|
||||
if client is None:
|
||||
raise ValueError("no rerank endpoint (base_url is empty)")
|
||||
return calibrate(client, model=model or base_url)
|
||||
try:
|
||||
return calibrate(client, model=model or base_url)
|
||||
finally:
|
||||
close = getattr(client, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
|
||||
|
||||
def _raw_scale(raw: list[float]) -> str:
|
||||
|
||||
@@ -1,63 +1,84 @@
|
||||
"""Resolve the runtime rerank client from configuration.
|
||||
"""Resolve an immutable runtime rerank lane from configuration.
|
||||
|
||||
Sole caller is ``ChatSession._resolve_rerank_client``: the reranker is the model
|
||||
Sole caller is ``ChatSession._resolve_rerank_lane``: the reranker is the model
|
||||
definition (capability ``supports_rerank``) selected via the Reranker role
|
||||
(``tools.reranker_alias``); there is no global endpoint fallback. The calibrate
|
||||
CLI (``admin.py``) and the calibrate endpoint (``_global_rerank_instruction`` in
|
||||
console/server.py) resolve the endpoint themselves via ``calibrate_model`` and
|
||||
only share the instruction precedence below, not this function.
|
||||
(``tools.reranker_alias``); there is no global endpoint fallback. Resolution
|
||||
creates a fresh frozen binding around a registry-owned shared runtime, never a
|
||||
session-owned client. Calibration remains an isolated one-shot path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from turnstone.core.rerank import RerankClient
|
||||
from turnstone.core.rerank import RerankLane
|
||||
|
||||
|
||||
def resolve_rerank_client_from(
|
||||
def _effective_snapshot(config_store: Any) -> tuple[int, dict[str, Any]]:
|
||||
snapshot = getattr(config_store, "effective_snapshot", None)
|
||||
if callable(snapshot):
|
||||
version, settings = snapshot()
|
||||
return int(version), dict(settings)
|
||||
# Compatibility for lightweight embedders/tests implementing only get().
|
||||
return int(getattr(config_store, "version", 0)), {
|
||||
"tools.reranker_alias": config_store.get("tools.reranker_alias"),
|
||||
"tools.rerank_instruction": config_store.get("tools.rerank_instruction"),
|
||||
}
|
||||
|
||||
|
||||
def _deactivate(registry: Any) -> None:
|
||||
deactivate = getattr(registry, "deactivate_rerank_runtime", None)
|
||||
if callable(deactivate):
|
||||
deactivate()
|
||||
|
||||
|
||||
def resolve_rerank_lane_from(
|
||||
config_store: Any | None,
|
||||
registry: Any | None,
|
||||
*,
|
||||
timeout: float,
|
||||
) -> RerankClient | None:
|
||||
"""Return a rerank client from the selected Reranker model, or ``None``.
|
||||
max_attempts: int = 3,
|
||||
) -> RerankLane | None:
|
||||
"""Return a coherent lane for the selected Reranker model, or ``None``.
|
||||
|
||||
The reranker is the model definition (capability ``supports_rerank``) picked
|
||||
via the Reranker role (``tools.reranker_alias``) — managed like every other
|
||||
model, its ``base_url`` the full Cohere/Jina-compatible /rerank endpoint.
|
||||
Returns ``None`` until such a model is selected (there is no bundled rerank
|
||||
endpoint and no global URL fallback).
|
||||
model, its ``base_url`` the full Cohere/Jina-compatible /rerank endpoint. A
|
||||
ConfigStore snapshot and registry generation witness keep multi-key settings
|
||||
and model binding coherent across concurrent reloads.
|
||||
"""
|
||||
if config_store is None or registry is None:
|
||||
return None
|
||||
|
||||
cs = config_store
|
||||
alias = str(cs.get("tools.reranker_alias") or "").strip()
|
||||
if not alias:
|
||||
return None
|
||||
try:
|
||||
cfg = registry.get_config(alias)
|
||||
except Exception:
|
||||
cfg = None
|
||||
if cfg is None or not cfg.base_url or not cfg.capabilities.get("supports_rerank"):
|
||||
return None
|
||||
|
||||
from turnstone.core.config import get_rerank_instruction
|
||||
from turnstone.core.rerank import resolve_rerank_client
|
||||
|
||||
# The query instruction is a global task knob (instruction-aware rerankers
|
||||
# like Qwen3); it applies to whichever reranker model is active. Stored admin
|
||||
# value wins, else config.toml/env -- the same precedence the calibrate CLI
|
||||
# (admin.py) and the calibrate endpoint (_global_rerank_instruction) use, so
|
||||
# the instruction used at calibration time matches the one used at runtime.
|
||||
instruction = str(cs.get("tools.rerank_instruction") or "").strip() or get_rerank_instruction()
|
||||
|
||||
return resolve_rerank_client(
|
||||
url=cfg.base_url,
|
||||
model=cfg.model or "",
|
||||
api_key=cfg.api_key,
|
||||
timeout=timeout,
|
||||
instruction=instruction,
|
||||
)
|
||||
for _attempt in range(max(1, max_attempts)):
|
||||
version, settings = _effective_snapshot(config_store)
|
||||
alias = str(settings.get("tools.reranker_alias") or "").strip()
|
||||
if not alias:
|
||||
_deactivate(registry)
|
||||
return None
|
||||
instruction = (
|
||||
str(settings.get("tools.rerank_instruction") or "").strip() or get_rerank_instruction()
|
||||
)
|
||||
try:
|
||||
lane = cast(
|
||||
"RerankLane",
|
||||
registry.resolve_rerank_lane(
|
||||
alias,
|
||||
instruction=instruction,
|
||||
config_version=version,
|
||||
),
|
||||
)
|
||||
except (KeyError, ValueError):
|
||||
_deactivate(registry)
|
||||
return None
|
||||
if (
|
||||
int(getattr(config_store, "version", version)) == version
|
||||
and int(getattr(registry, "generation", lane.registry_generation))
|
||||
== lane.registry_generation
|
||||
):
|
||||
return lane
|
||||
# A continuously changing configuration is not a safe binding. Retrieval
|
||||
# callers preserve native order and retry resolution on their next use.
|
||||
return None
|
||||
|
||||
+136
-42
@@ -278,7 +278,7 @@ if TYPE_CHECKING:
|
||||
)
|
||||
from turnstone.core.output_guard import OutputAssessment
|
||||
from turnstone.core.output_guard_judge import OutputGuardJudge, OutputJudgeVerdict
|
||||
from turnstone.core.rerank import RerankClient, Reranker
|
||||
from turnstone.core.rerank import Reranker, RerankHit, RerankLane
|
||||
from turnstone.core.storage import ForkCloneSnapshot
|
||||
from turnstone.core.web_search import WebSearchClient
|
||||
|
||||
@@ -3505,28 +3505,80 @@ class ChatSession:
|
||||
timeout=self.tool_timeout,
|
||||
)
|
||||
|
||||
def _resolve_rerank_client(self) -> RerankClient | None:
|
||||
"""Return a rerank client, or None when reranking is unconfigured.
|
||||
def _resolve_rerank_lane(self) -> RerankLane | None:
|
||||
"""Return a fresh rerank lane, or None when reranking is unconfigured.
|
||||
|
||||
The reranker is a **model definition** (capability ``supports_rerank``)
|
||||
selected via the Reranker role (``tools.reranker_alias``); its base_url
|
||||
is the full /rerank endpoint. There is no bundled rerank endpoint and no
|
||||
global URL fallback, so reranking stays disabled until such a model is
|
||||
selected.
|
||||
global URL fallback. The registry owns the shared runtime; callers must
|
||||
resolve per batch rather than retain this immutable binding.
|
||||
"""
|
||||
from turnstone.core.rerank_config import resolve_rerank_client_from
|
||||
from turnstone.core.rerank_config import resolve_rerank_lane_from
|
||||
|
||||
return resolve_rerank_client_from(
|
||||
return resolve_rerank_lane_from(
|
||||
getattr(self, "_config_store", None),
|
||||
getattr(self, "_registry", None),
|
||||
timeout=min(self.tool_timeout, _RERANK_TIMEOUT_CAP_S),
|
||||
)
|
||||
|
||||
def _rerank_hits(
|
||||
self,
|
||||
query: str,
|
||||
documents: list[str],
|
||||
*,
|
||||
origin_generation: int = 0,
|
||||
) -> list[RerankHit] | None:
|
||||
"""Resolve and dispatch one batch, preserving its cancellation owner.
|
||||
|
||||
``None`` means no live reranker is configured; callers return identity
|
||||
order so long-lived BM25/tool-search closures can observe a later role
|
||||
selection without being rebuilt. Endpoint and circuit failures raise
|
||||
through the existing retrieval fallback seams.
|
||||
"""
|
||||
lane = self._resolve_rerank_lane()
|
||||
if lane is None:
|
||||
return None
|
||||
|
||||
from turnstone.core.deadline import DeadlineCancelledError
|
||||
from turnstone.core.rerank import rerank
|
||||
|
||||
task_scope = _active_task_agent_cancel_scope.get()
|
||||
owner_generation = (
|
||||
origin_generation
|
||||
or _active_tool_origin_generation.get()
|
||||
or _active_commit_origin_generation.get()
|
||||
)
|
||||
if task_scope is not None:
|
||||
cancel_ref: Any = task_scope.cancel_ref
|
||||
elif hasattr(self, "_cancel_event"):
|
||||
cancel_ref = _CancelRef(self, owner_generation)
|
||||
else: # lightweight direct seam tests/embedders
|
||||
cancel_ref = None
|
||||
|
||||
try:
|
||||
return rerank(
|
||||
lane,
|
||||
query,
|
||||
documents,
|
||||
timeout=min(float(getattr(self, "tool_timeout", 30.0)), _RERANK_TIMEOUT_CAP_S),
|
||||
cancel_ref=cancel_ref,
|
||||
)
|
||||
except DeadlineCancelledError:
|
||||
# Retrieval fallbacks catch Exception. Translate the deadline seam
|
||||
# back into GenerationCancelled (a BaseException) before they can
|
||||
# mistake Stop/supersession for an endpoint outage.
|
||||
if task_scope is not None:
|
||||
task_scope.check()
|
||||
check_cancelled = getattr(self, "_check_cancelled", None)
|
||||
if callable(check_cancelled):
|
||||
check_cancelled(owner_generation)
|
||||
raise
|
||||
|
||||
def _rerank_enabled_for(self, tool: str) -> bool:
|
||||
"""Whether reranking is enabled for ``tool`` (currently: 'web_search', 'bm25').
|
||||
|
||||
The per-tool toggles default on; the operative gate is whether an
|
||||
endpoint is configured (``_resolve_rerank_client`` returns None when
|
||||
endpoint is configured (``_resolve_rerank_lane`` returns None when
|
||||
not). A bare CLI without a ConfigStore inherits the on-by-default toggle.
|
||||
"""
|
||||
cs = getattr(self, "_config_store", None)
|
||||
@@ -3534,25 +3586,30 @@ class ChatSession:
|
||||
return bool(cs.get(f"tools.rerank_{tool}"))
|
||||
return True
|
||||
|
||||
def _web_search_reranker(self) -> Reranker | None:
|
||||
"""Build a web_search reranker callable, or None when disabled.
|
||||
def _web_search_reranker(self) -> Reranker:
|
||||
"""Build a live-resolving web_search reranker callable.
|
||||
|
||||
Returns a ``(query, docs) -> ranked indices`` adapter over the configured
|
||||
rerank endpoint; None when reranking is off or no endpoint is set.
|
||||
The closure captures no concrete lane/runtime. Config changes are read
|
||||
at each invocation; disabled/unconfigured state is identity order.
|
||||
"""
|
||||
if not self._rerank_enabled_for("web_search"):
|
||||
return None
|
||||
rc = self._resolve_rerank_client()
|
||||
if rc is None:
|
||||
return None
|
||||
|
||||
def _rank(query: str, docs: list[str]) -> list[int]:
|
||||
return [hit.index for hit in rc.rerank(query, docs)]
|
||||
if not self._rerank_enabled_for("web_search"):
|
||||
return list(range(len(docs)))
|
||||
hits = self._rerank_hits(query, docs)
|
||||
if hits is None:
|
||||
return list(range(len(docs)))
|
||||
return [hit.index for hit in hits]
|
||||
|
||||
return _rank
|
||||
|
||||
def _bm25_reranker(self, threshold: float = 0.0) -> Reranker | None:
|
||||
"""Build a BM25 reranker callable, or None when disabled.
|
||||
def _bm25_reranker(
|
||||
self,
|
||||
threshold: float = 0.0,
|
||||
*,
|
||||
origin_generation: int = 0,
|
||||
) -> Reranker:
|
||||
"""Build a live-resolving BM25 reranker callable.
|
||||
|
||||
Mirrors ``_web_search_reranker``. ``threshold`` is a relevance FLOOR
|
||||
applied in this closure (where scores still exist); the BM25Index seam
|
||||
@@ -3563,24 +3620,22 @@ class ChatSession:
|
||||
threshold. Only memory-pointer relevance filtering passes a configured threshold;
|
||||
reactive surfaces pass 0.
|
||||
"""
|
||||
if not self._rerank_enabled_for("bm25"):
|
||||
return None
|
||||
rc = self._resolve_rerank_client()
|
||||
if rc is None:
|
||||
return None
|
||||
|
||||
from turnstone.core.rerank import RerankError, normalize_scores
|
||||
|
||||
def _rank(query: str, docs: list[str]) -> list[int]:
|
||||
hits = rc.rerank(query, docs)
|
||||
if not self._rerank_enabled_for("bm25"):
|
||||
return list(range(len(docs)))
|
||||
hits = self._rerank_hits(
|
||||
query,
|
||||
docs,
|
||||
origin_generation=origin_generation,
|
||||
)
|
||||
if hits is None:
|
||||
return list(range(len(docs)))
|
||||
if docs and not hits:
|
||||
# A conforming reranker scores every document; an empty result
|
||||
# for non-empty input means the endpoint response was
|
||||
# unparseable (rc.rerank -> _parse_hits returns [] without
|
||||
# raising). That is an endpoint FAILURE -- a discrete branch
|
||||
# from the relevance floor below -- so raise and let BM25Index
|
||||
# fall back to BM25 order in BOTH modes, instead of the
|
||||
# filter-mode floor honoring it as "nothing relevant".
|
||||
# The managed dispatcher enforces this too. Keep the adapter
|
||||
# fail-closed for alternate protocol implementations and
|
||||
# lightweight test embedders that provide the seam directly.
|
||||
raise RerankError("rerank endpoint returned no scores for non-empty input")
|
||||
# Normalise to a 0-1 relevance space (sigmoid for logit endpoints) so
|
||||
# ``threshold`` means the same on every reranker. Sigmoid is monotonic
|
||||
@@ -11204,7 +11259,13 @@ class ChatSession:
|
||||
# below rejects every old mutation. The participant lookup is cached so
|
||||
# the in-lock apply path performs no storage access.
|
||||
planned_user_turn = (
|
||||
self._plan_genuine_user_turn(user_input, turn_principal_id) if not from_wake else None
|
||||
self._plan_genuine_user_turn(
|
||||
user_input,
|
||||
turn_principal_id,
|
||||
origin_generation=my_generation,
|
||||
)
|
||||
if not from_wake
|
||||
else None
|
||||
)
|
||||
shared_state_plan = self._plan_shared_state()
|
||||
participant_name: str | None = None
|
||||
@@ -11885,7 +11946,9 @@ class ChatSession:
|
||||
# must not let this retired worker pop a successor's queue
|
||||
# or repaint its state as idle.
|
||||
flushed_out: list[bool] = []
|
||||
queued_flush = self._prepare_queued_flush()
|
||||
queued_flush = self._prepare_queued_flush(
|
||||
origin_generation=my_generation,
|
||||
)
|
||||
|
||||
def _drain_or_idle(
|
||||
durable: list[Callable[[], None]],
|
||||
@@ -12396,6 +12459,7 @@ class ChatSession:
|
||||
queued_flush = self._prepare_queued_flush(
|
||||
user_feedback or "",
|
||||
include_queue=False,
|
||||
origin_generation=my_generation,
|
||||
)
|
||||
fold_tool_batch = functools.partial(
|
||||
_fold_tool_batch,
|
||||
@@ -12468,7 +12532,10 @@ class ChatSession:
|
||||
)
|
||||
raise
|
||||
except GenerationCancelled:
|
||||
queued_flush = self._prepare_queued_flush()
|
||||
queued_flush = self._prepare_queued_flush(
|
||||
origin_generation=my_generation,
|
||||
allow_rerank=False,
|
||||
)
|
||||
|
||||
def _finalize_cancelled_generation(
|
||||
durable: list[Callable[[], None]],
|
||||
@@ -12575,7 +12642,10 @@ class ChatSession:
|
||||
# Do NOT re-raise — return normally so server worker thread
|
||||
# completes cleanly.
|
||||
except KeyboardInterrupt as exc:
|
||||
queued_flush = self._prepare_queued_flush()
|
||||
queued_flush = self._prepare_queued_flush(
|
||||
origin_generation=my_generation,
|
||||
allow_rerank=False,
|
||||
)
|
||||
|
||||
def _finalize_interrupted_generation(
|
||||
error: BaseException,
|
||||
@@ -12601,7 +12671,10 @@ class ChatSession:
|
||||
)
|
||||
raise
|
||||
except Exception as exc:
|
||||
queued_flush = self._prepare_queued_flush()
|
||||
queued_flush = self._prepare_queued_flush(
|
||||
origin_generation=my_generation,
|
||||
allow_rerank=False,
|
||||
)
|
||||
|
||||
# Orphan gate: a superseded thread's stream death can escape
|
||||
# cancel conversion (the successor replaced the cancel event
|
||||
@@ -17099,6 +17172,8 @@ class ChatSession:
|
||||
prefix: str = "",
|
||||
*,
|
||||
include_queue: bool = True,
|
||||
origin_generation: int = 0,
|
||||
allow_rerank: bool = True,
|
||||
) -> _QueuedFlushPlan:
|
||||
"""Snapshot and rank one prospective queued USER turn outside locks."""
|
||||
worker_claim = current_worker_claim(self)
|
||||
@@ -17133,7 +17208,12 @@ class ChatSession:
|
||||
return _QueuedFlushPlan(
|
||||
prefix=prefix,
|
||||
items=items,
|
||||
turn=self._plan_genuine_user_turn(content, principal),
|
||||
turn=self._plan_genuine_user_turn(
|
||||
content,
|
||||
principal,
|
||||
origin_generation=origin_generation,
|
||||
allow_rerank=allow_rerank,
|
||||
),
|
||||
)
|
||||
|
||||
def _apply_queued_flush(
|
||||
@@ -19456,6 +19536,8 @@ class ChatSession:
|
||||
user_message: str,
|
||||
*,
|
||||
access: _MemoryAccess,
|
||||
origin_generation: int = 0,
|
||||
allow_rerank: bool = True,
|
||||
) -> str:
|
||||
"""Build a live metadata-only pointer for one genuine user turn."""
|
||||
if not user_message.strip() or not self._nudges_enabled("memory_pointer"):
|
||||
@@ -19470,7 +19552,14 @@ class ChatSession:
|
||||
rows,
|
||||
user_message,
|
||||
k=self._mem_cfg.relevance_k,
|
||||
reranker=self._bm25_reranker(threshold),
|
||||
reranker=(
|
||||
self._bm25_reranker(
|
||||
threshold,
|
||||
origin_generation=origin_generation,
|
||||
)
|
||||
if allow_rerank
|
||||
else None
|
||||
),
|
||||
rerank_filters=threshold > 0,
|
||||
)
|
||||
except Exception:
|
||||
@@ -19482,6 +19571,9 @@ class ChatSession:
|
||||
self,
|
||||
content: str,
|
||||
principal_id: str,
|
||||
*,
|
||||
origin_generation: int = 0,
|
||||
allow_rerank: bool = True,
|
||||
) -> _GenuineUserTurnPlan:
|
||||
"""Plan all metadata derived from an admitted human turn."""
|
||||
principal = principal_id.strip()
|
||||
@@ -19491,6 +19583,8 @@ class ChatSession:
|
||||
memory_pointer=self._plan_memory_pointer(
|
||||
content,
|
||||
access=self._memory_access(principal),
|
||||
origin_generation=origin_generation,
|
||||
allow_rerank=allow_rerank,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user