From bfc05bf3b79976c52439008f2b6dd5c4c5a53b0e Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:54:31 +0800 Subject: [PATCH] feat(decision): allow hosts to inject owned JEV model instances (#652) * feat(decision): support host-owned model injection * fix(decision): finish model cleanup before propagating shutdown cancellation --- docs/en/context_compaction.md | 17 ++ docs/zh/context_compaction.md | 15 ++ .../as_decision_model/agentscope.py | 52 +++++- tests/unit/test_decision_model.py | 159 ++++++++++++++++++ 4 files changed, 235 insertions(+), 8 deletions(-) diff --git a/docs/en/context_compaction.md b/docs/en/context_compaction.md index 17ca97b3..c6677668 100644 --- a/docs/en/context_compaction.md +++ b/docs/en/context_compaction.md @@ -159,3 +159,20 @@ the entire gain cannot be attributed to Jev ranking itself. Agent continuation s [Complete experiment: design, results, expected answers and reproduction](https://github.com/jinliyl/ReMe/blob/74e1251ddcd212205788f21ccd6b5ebd82ade7c3/benchmark/compaction/README.md) + +## Host-owned decision models + +Embedded hosts can reuse their existing AgentScope JEV instance, just as they inject an LLM: + +```python +# app contains an as_decision_model.compaction component with backend: agentscope. +await app.update_component("as_decision_model", "compaction", model=host_jev) +``` + +The YAML `model` value remains the provider model name. The runtime `model` property accepts a callable instance +and does not construct a client on access. Inject before the first request; ReMe validates and converts JSON +questions normally, but does not require separate ReMe credentials or close the borrowed client. The host must +keep the instance alive until ReMe has finished its calls, then close it itself. Close the component before +replacing an internally constructed model or any model with active requests. Shutdown waits for active calls and +owned-client cleanup even when cancelled, then propagates cancellation. Component shutdown detaches the instance; +inject again before restarting a borrowed-model component. Standalone credential-based behavior is unchanged. diff --git a/docs/zh/context_compaction.md b/docs/zh/context_compaction.md index 4de6352a..33accb2e 100644 --- a/docs/zh/context_compaction.md +++ b/docs/zh/context_compaction.md @@ -147,3 +147,18 @@ agent 继续执行的成功率尚未验证。 [完整实验：设计、结果、标准答案与复现](https://github.com/jinliyl/ReMe/blob/74e1251ddcd212205788f21ccd6b5ebd82ade7c3/benchmark/compaction/README_ZH.md) + +## 注入宿主管理的决策模型 + +嵌入式宿主可以复用已有的 AgentScope JEV 实例，与注入 LLM 的方式一致： + +```python +# app 已配置 backend: agentscope 的 as_decision_model.compaction 组件。 +await app.update_component("as_decision_model", "compaction", model=host_jev) +``` + +YAML 中的 `model` 仍表示供应商模型名称；运行时 `model` 属性接收可调用的模型实例，读取属性不会创建客户端。 +在首次请求之前注入。ReMe 继续校验并转换 JSON 问题，但不再要求额外的 ReMe 凭据，也不会关闭借用的客户端。 +宿主应等待 ReMe 调用结束后自行关闭实例。替换 ReMe 自建模型或存在进行中请求的实例之前，须先关闭组件。 +关闭过程即使被取消，也会等待活动请求和自建客户端清理完成，再传播取消信号。组件关闭时会解除实例引用； +重新启动使用外部模型的组件前需要再次注入。独立运行时按凭据创建模型的行为保持不变。 diff --git a/reme/components/as_decision_model/agentscope.py b/reme/components/as_decision_model/agentscope.py index 2e8ff025..640e3a93 100644 --- a/reme/components/as_decision_model/agentscope.py +++ b/reme/components/as_decision_model/agentscope.py @@ -31,8 +31,8 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): """Wrap AgentScope's TypeSafe decision model with lazy provider loading. ``questions`` accepts AgentScope question objects or JSON dictionaries with - a ``type`` discriminator (``binary``, ``choice``, or ``score``). Only this - component's configuration supplies credentials and transport options. + a ``type`` discriminator (``binary``, ``choice``, or ``score``). Hosts may inject a borrowed model + through ``model``; otherwise this component constructs and owns its client. """ def __init__( @@ -53,10 +53,31 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): retry_delay=retry_delay, ) self._model: "ClassifierModelBase | None" = None + self._owns_model = False self._active_calls = 0 self._idle = asyncio.Event() self._idle.set() + @property + def model(self) -> "ClassifierModelBase | None": + """Return the current instance without constructing a provider client.""" + return self._model + + @model.setter + def model(self, value: "ClassifierModelBase | None") -> None: + """Borrow a host-owned instance, including via Application.update_component. + + Close the component before replacing an internally owned client. An + active request must finish before any replacement or detachment. + """ + if value is self._model: + return + if value is not None and not callable(value): + raise TypeError("An injected decision model must be callable.") + if self._active_calls or self._owns_model: + raise RuntimeError("Close the decision component before replacing its active or owned model.") + self._model = value + @staticmethod def _load_decision_model(): """Load the optional AgentScope adapter on the first decision request.""" @@ -102,6 +123,7 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): credential=credential_cls(**credential), **self.options.model_dump(exclude_none=True), ) + self._owns_model = True except ImportError as exc: raise ImportError( "Decision model inference requires a compatible typesafe-sdk. " @@ -120,10 +142,24 @@ class AgentScopeAsDecisionModel(BaseAsDecisionModel): self._idle.set() async def _close(self) -> None: + """Finish draining and cleanup before propagating shutdown cancellation.""" + cleanup = asyncio.create_task(self._release_model()) + cancellation = None + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError as exc: + # Keep the lifecycle lock until cleanup finishes, even after + # repeated cancellation. No cleanup task outlives close(). + cancellation = exc + cleanup.result() + if cancellation is not None: + raise cancellation + + async def _release_model(self) -> None: """Release the owned TypeSafe client and allow a later fresh start.""" - try: - await self._idle.wait() - finally: - model, self._model = self._model, None - if model is not None: - await model.client.aclose() + await self._idle.wait() + model, owned = self._model, self._owns_model + self._model, self._owns_model = None, False + if model is not None and owned: + await model.client.aclose() diff --git a/tests/unit/test_decision_model.py b/tests/unit/test_decision_model.py index 1ec36779..1a4c3beb 100644 --- a/tests/unit/test_decision_model.py +++ b/tests/unit/test_decision_model.py @@ -238,6 +238,88 @@ def test_failed_or_cancelled_call_does_not_block_close(provider, failure): asyncio.run(run()) +@pytest.mark.parametrize("borrowed", [False, True]) +def test_cancelled_close_drains_requests_and_releases_model(provider, borrowed): + """Shutdown defers cancellation until active calls and model cleanup finish.""" + _, model = provider + + async def run(): + started, release = asyncio.Event(), asyncio.Event() + + async def reply(**_kwargs): + started.set() + await release.wait() + + model.side_effect = reply + component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"}) + if borrowed: + component.model = model + call = asyncio.create_task(component(state="hello", questions={"q": {"type": "binary"}})) + await asyncio.wait_for(started.wait(), 2) + close = asyncio.create_task(component.close()) + await asyncio.sleep(0) + try: + for _ in range(2): + close.cancel() + await asyncio.sleep(0) + assert not close.done() + assert component.is_started + assert component.model is model + model.client.aclose.assert_not_awaited() + release.set() + await asyncio.wait_for(call, 2) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(close, 2) + assert component.model is None + assert not component.is_started + assert not component._owns_model + await component.close() + assert model.client.aclose.await_count == (0 if borrowed else 1) + finally: + release.set() + await asyncio.gather(call, close, return_exceptions=True) + + asyncio.run(run()) + + +def test_cancelled_close_finishes_owned_client_cleanup(provider): + """Repeated shutdown cancellation cannot interrupt the client's own close.""" + _, model = provider + + async def run(): + closing, release, finished = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def close_client(): + closing.set() + await release.wait() + finished.set() + + model.client.aclose.side_effect = close_client + component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"}) + await component(state="hello", questions={"q": {"type": "binary"}}) + close = asyncio.create_task(component.close()) + await asyncio.wait_for(closing.wait(), 2) + try: + for _ in range(2): + close.cancel() + await asyncio.sleep(0) + assert not close.done() + assert component.is_started + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(close, 2) + assert finished.is_set() + assert component.model is None + assert not component.is_started + await component.close() + model.client.aclose.assert_awaited_once() + finally: + release.set() + await asyncio.gather(close, return_exceptions=True) + + asyncio.run(run()) + + def test_client_close_error_clears_cached_model(provider): """A failed close cannot leave a stale client selected on the next start.""" factory, model = provider @@ -350,3 +432,80 @@ def test_adapter_uses_sdk_default_or_explicit_model(monkeypatch, model_name): model.client.aclose.assert_awaited_once() asyncio.run(run()) + + +def test_injected_model_is_borrowed_and_accepts_json_questions(provider, tmp_path): + """Application injection shares the host instance without owning its client.""" + from reme.application import Application + + factory, model = provider + + async def run(): + app = Application( + workspace_dir=str(tmp_path), + enable_logo=False, + log_to_console=False, + log_to_file=False, + service={"backend": "http"}, + components={"as_decision_model": {"compaction": {"backend": "agentscope"}}}, + ) + component = await app.update_component("as_decision_model", "compaction", model=model) + assert component.model is model + async with app: + await component(state="hello", questions={"q": {"type": "binary"}}) + factory.assert_not_called() + model.assert_awaited_once() + model.client.aclose.assert_not_awaited() + assert component.model is None + + asyncio.run(run()) + + +def test_model_replacement_guards_owned_clients(provider): + """Replacing an owned client requires an explicit close to prevent leaks.""" + _, model = provider + + async def run(): + component = AgentScopeAsDecisionModel(credential={"api_key": "test-key"}) + await component(state="hello", questions={"q": {"type": "binary"}}) + component.model = model + with pytest.raises(RuntimeError, match="Close"): + component.model = None + await component.close() + component.model = model + await component.close() + # The borrowed instance is never closed a second time. + model.client.aclose.assert_awaited_once() + with pytest.raises(TypeError, match="callable"): + component.model = "model-name" + + asyncio.run(run()) + + +def test_borrowed_model_cannot_be_replaced_while_in_use(provider): + """Host ownership still obeys the component's in-flight request boundary.""" + _, model = provider + + async def run(): + started, release = asyncio.Event(), asyncio.Event() + + async def reply(**_kwargs): + started.set() + await release.wait() + + model.side_effect = reply + component = AgentScopeAsDecisionModel() + component.model = model + task = asyncio.create_task(component(state="hello", questions={"q": {"type": "binary"}})) + await asyncio.wait_for(started.wait(), 2) + with pytest.raises(RuntimeError, match="Close"): + component.model = None + close = asyncio.create_task(component.close()) + await asyncio.sleep(0) + assert not close.done() + release.set() + await task + await close + model.client.aclose.assert_not_awaited() + + asyncio.run(run())