mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
feat(extensions): grant bounded host model invocation (#5796)
* feat(extensions): grant bounded host model invocation * fix(extensions): retain model budgets and isolate schema validation * fix(extensions): normalize provider-origin model cancellation --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
yang rui
GitHub
parent
b4a60fbfe2
commit
c814943215
@ -1405,6 +1405,15 @@ returns `None` when the host does not support the capability; the required helpe
|
||||
these exceptions to HTTP 503 and 403 respectively; resolver failures never fall back to
|
||||
the global service reader.
|
||||
|
||||
Extension services can optionally receive a [host model invoker](backend/docs/extension-model-invocation.md)
|
||||
(extension API 0.2.4+). Operators explicitly map allowed logical roles under
|
||||
`plugins[].host_access.model_invocation`; the host handles credentials, concurrency,
|
||||
timeouts, and schema-validated output. Admission is bounded; timed-out provider work
|
||||
keeps its concurrency slot until completion, and schema validation runs in terminable
|
||||
child processes. A provider's own cancellation becomes `ModelInvocationFailed`,
|
||||
while cancellation of the calling extension task still propagates normally.
|
||||
Without a grant the capability is `None`.
|
||||
|
||||
Plugin order is deterministic, per-plugin configuration is passed to `install()`, and
|
||||
`required: true` makes load failure abort startup; otherwise failures are reported and
|
||||
skipped. `enabled: false` skips resolution and import. The manager preserves the extension's
|
||||
|
||||
155
backend/docs/extension-model-invocation.md
Normal file
155
backend/docs/extension-model-invocation.md
Normal file
@ -0,0 +1,155 @@
|
||||
# Host model invocation for extensions
|
||||
|
||||
Extension API **0.2.4** adds an optional `ExtensionRuntimeDeps.model_invoker`.
|
||||
An operator-authorized service can make an asynchronous, non-streaming text call
|
||||
using DeerFlow's configured models. The host constructs the provider client and
|
||||
returns plain data; the extension needs no provider credentials or LangChain dependency.
|
||||
|
||||
## Grant and route logical roles
|
||||
|
||||
Add `host_access` to the relevant entry in the operator-owned `config.yaml`:
|
||||
|
||||
```yaml
|
||||
plugins:
|
||||
- name: classifier
|
||||
use: my_classifier:install
|
||||
host_access:
|
||||
model_invocation:
|
||||
roles:
|
||||
default: my-configured-model
|
||||
fast: my-small-model
|
||||
max_concurrency: 2
|
||||
timeout_seconds: 60
|
||||
max_input_chars: 262144
|
||||
max_output_chars: 65536
|
||||
```
|
||||
|
||||
The mapping values are names from `models:`. Role names are deployment-defined;
|
||||
an omitted request role selects `default`, which must be explicitly mapped.
|
||||
An extension cannot request an arbitrary configured model name. No grant (or
|
||||
`model_invocation: null`) means `deps.model_invoker is None`, preserving existing
|
||||
services. `host_access` is separate from the extension-private `config` passed to
|
||||
`install()` and cannot be edited through `extensions_config.json`.
|
||||
|
||||
Each successful installation owns its grant and concurrency budget. Multiple
|
||||
services from that installation share the budget; another installation using
|
||||
the same Python entry point does not inherit its grant. Grants and model
|
||||
configuration are bound at service startup. As with all `plugins:` changes,
|
||||
changing the grant, source, or extension model configuration requires a Gateway
|
||||
restart. Startup failure revokes that service's handle. Shutdown revokes retained
|
||||
handles and cancels callers before invoking its `stop()` method. Already-running
|
||||
provider work retains its budget until it actually finishes; its result is discarded.
|
||||
|
||||
This capability follows the trusted extension model: installed Python code already
|
||||
runs with Gateway privileges. It is not a sandbox or a per-user authorization API.
|
||||
A service exposing an HTTP route must authorize callers and their input itself.
|
||||
|
||||
## Call from a service
|
||||
|
||||
```python
|
||||
from deerflow_extension_api import (
|
||||
ModelInvocationRequest,
|
||||
ModelMessage,
|
||||
extension,
|
||||
)
|
||||
|
||||
|
||||
class Classifier:
|
||||
async def start(self, deps):
|
||||
self.invoker = deps.model_invoker
|
||||
|
||||
async def stop(self):
|
||||
pass
|
||||
|
||||
async def classify(self, text):
|
||||
if self.invoker is None:
|
||||
raise RuntimeError("Operator has not granted model invocation")
|
||||
result = await self.invoker.invoke(ModelInvocationRequest(
|
||||
messages=[ModelMessage("user", text)],
|
||||
purpose="classification",
|
||||
response_schema={
|
||||
"type": "object",
|
||||
"properties": {"label": {"enum": ["positive", "negative"]}},
|
||||
"required": ["label"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
))
|
||||
return result.structured_output["label"]
|
||||
|
||||
|
||||
@extension(api="0.2.4")
|
||||
def install(registry, config):
|
||||
registry.service(Classifier())
|
||||
```
|
||||
|
||||
Invoke from the service's event loop. Messages support `system`, `user`, and
|
||||
`assistant` text, with 1–256 messages per call. `purpose` is an optional 1–128
|
||||
character tracing label, not a prompt or an authorization selector. Do not put
|
||||
secrets in it. The caller can shorten `timeout_seconds`, but cannot extend the
|
||||
host timeout. Queueing, provider construction, schema-worker startup, and validation
|
||||
count toward the same deadline. Cancellation propagates as `asyncio.CancelledError`.
|
||||
Timeout and cancellation stop waiting promptly, but cannot terminate synchronous
|
||||
provider threads. Both synchronous and asynchronous provider calls therefore keep
|
||||
their concurrency and admission slots until the actual operation finishes. No
|
||||
additional provider request starts if construction finishes after its caller left.
|
||||
Configure provider-side timeouts as well: a stuck provider keeps its slot occupied.
|
||||
Stopping a service does not wait for those provider operations to finish.
|
||||
|
||||
`response_schema` accepts inline JSON Schema Draft 2020-12 object schemas.
|
||||
References (`$ref`, `$dynamicRef`, `$recursiveRef`) and other schema dialects are
|
||||
rejected before dispatch; schema validation never fetches network resources.
|
||||
The host adds a JSON-only instruction and validates the returned text locally,
|
||||
so the contract does not depend on a provider's native structured-output feature.
|
||||
Schema checking and response validation each run in a short-lived isolated Python
|
||||
process using the host interpreter. This adds process startup overhead to structured
|
||||
calls. Expensive regexes or schema combinations cannot block the Gateway event loop
|
||||
or its GIL. On timeout/cancellation the host kills and reaps the child before
|
||||
releasing admission (pipe cancellation is polled every 50 ms). No schema-worker
|
||||
process is used for plain text calls.
|
||||
Malformed JSON, non-finite numbers, and schema violations are explicit failures;
|
||||
no partially populated success is returned. Batching, row-ID reconciliation,
|
||||
file handling, and workflow-specific checks belong to the extension.
|
||||
|
||||
Successful results contain `content`, optional validated `structured_output`,
|
||||
the resolved configured model name, and optional `ModelUsage` token counts.
|
||||
Provider response metadata, framework messages, and exceptions are not exposed.
|
||||
The standard model factory retains provider adaptation, configured RPM admission,
|
||||
and tracing callbacks; invocation metadata attributes calls to their extension
|
||||
source, logical role, and purpose.
|
||||
|
||||
## Failures and bounds
|
||||
|
||||
| Exception | Meaning |
|
||||
| --- | --- |
|
||||
| `ModelInvocationUnavailable` | Unsupported host, stopped capability, wrong event loop, or missing configured model |
|
||||
| `ModelInvocationUnauthorized` | The requested logical role is not granted |
|
||||
| `ModelInvocationFailed` | Invalid request, provider failure, timeout, unsupported response, or host limit exceeded |
|
||||
| `ModelOutputValidationError` | Malformed JSON or schema mismatch; a subclass of `ModelInvocationFailed` |
|
||||
|
||||
All these derive from `ModelInvocationError`. Provider failures use a normalized
|
||||
message without the original exception chain. Host deadline expiry reports
|
||||
`Model invocation timed out`; a provider's own `TimeoutError` reports
|
||||
`Model provider timed out`. A provider task that cancels itself reports
|
||||
`ModelInvocationFailed("Model provider cancelled")`; it does not cancel the
|
||||
extension's calling task. Normal caller cancellation remains
|
||||
`asyncio.CancelledError`, rather than being converted to an ordinary failure.
|
||||
|
||||
Concurrency defaults to 2 (range 1–64), timeout to 60 seconds (maximum 600), and
|
||||
input/output character limits to 262144/65536 (maximum 1048576 each). At most
|
||||
`2 * max_concurrency` calls are admitted per installation, including running,
|
||||
queued, and abandoned-but-still-running provider calls. Excess calls fail immediately
|
||||
with `ModelInvocationFailed("Model invocation capacity exceeded")`, before payload
|
||||
conversion or validation. Schema workers share the execution slots, so at most
|
||||
`max_concurrency` children can run per installation. Input
|
||||
counting includes the schema instruction. Output bounds apply to returned text
|
||||
after generation; configure the underlying model's token limits to bound provider
|
||||
generation costs. This is a per-installation, per-Gateway-process budget, not a
|
||||
distributed quota. Provider SDK retries retain their configured behavior.
|
||||
|
||||
Tool calls, streaming, agent execution, and conversation persistence are outside
|
||||
this capability. Jev/TypeSafe typed-decision workflows remain extension-owned;
|
||||
`model_invoker` is the chat-style path for host-configured models.
|
||||
|
||||
The offline tests in `backend/tests/test_extension_model_invocation.py` cover the
|
||||
classification success and failure cases, source isolation, shared admission,
|
||||
timeouts, cancellation, limits, and lifecycle revocation without real API calls.
|
||||
@ -43,6 +43,18 @@ from deerflow_extension_api.contracts import (
|
||||
TaskOutcome,
|
||||
extension,
|
||||
)
|
||||
from deerflow_extension_api.model_invocation import (
|
||||
ModelInvocationError,
|
||||
ModelInvocationFailed,
|
||||
ModelInvocationRequest,
|
||||
ModelInvocationResult,
|
||||
ModelInvocationUnauthorized,
|
||||
ModelInvocationUnavailable,
|
||||
ModelInvoker,
|
||||
ModelMessage,
|
||||
ModelOutputValidationError,
|
||||
ModelUsage,
|
||||
)
|
||||
from deerflow_extension_api.placement import (
|
||||
AgentBuildContext,
|
||||
AgentScope,
|
||||
@ -89,6 +101,16 @@ from deerflow_extension_api.state import ExtensionData
|
||||
API_VERSION = "0.2.4"
|
||||
|
||||
__all__ = [
|
||||
"ModelInvocationError",
|
||||
"ModelInvocationFailed",
|
||||
"ModelInvocationRequest",
|
||||
"ModelInvocationResult",
|
||||
"ModelInvocationUnauthorized",
|
||||
"ModelInvocationUnavailable",
|
||||
"ModelInvoker",
|
||||
"ModelMessage",
|
||||
"ModelOutputValidationError",
|
||||
"ModelUsage",
|
||||
"ActionContext",
|
||||
"BackendAction",
|
||||
"BrowserAssets",
|
||||
|
||||
@ -20,6 +20,7 @@ from deerflow_extension_api.state import ExtensionData
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from deerflow_extension_api.assembly import AgentAssemblyObserver
|
||||
from deerflow_extension_api.compaction import ContextCompactionObserver
|
||||
from deerflow_extension_api.model_invocation import ModelInvoker
|
||||
from deerflow_extension_api.placement import AgentBuildContext, MiddlewarePlacement
|
||||
from deerflow_extension_api.run_evidence import RunEvidenceReader
|
||||
|
||||
@ -167,6 +168,7 @@ class ExtensionRuntimeDeps:
|
||||
policy: HostPolicySnapshot = field(default_factory=HostPolicySnapshot)
|
||||
session_factory: Any | None = None
|
||||
run_evidence_reader: RunEvidenceReader | None = None
|
||||
model_invoker: ModelInvoker | None = None
|
||||
|
||||
|
||||
class ExtensionService(Protocol):
|
||||
|
||||
@ -0,0 +1,70 @@
|
||||
"""Provider-neutral, non-streaming text model calls granted by the host."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, Protocol
|
||||
|
||||
|
||||
class ModelInvocationError(RuntimeError):
|
||||
"""Normalized host error; contains no provider exception or credentials."""
|
||||
|
||||
|
||||
class ModelInvocationUnavailable(ModelInvocationError):
|
||||
"""The capability has stopped or its configured model is unavailable."""
|
||||
|
||||
|
||||
class ModelInvocationUnauthorized(ModelInvocationError):
|
||||
"""The requested logical role is not granted to this installation."""
|
||||
|
||||
|
||||
class ModelInvocationFailed(ModelInvocationError):
|
||||
"""The invocation failed, exceeded a host limit or timed out."""
|
||||
|
||||
|
||||
class ModelOutputValidationError(ModelInvocationFailed):
|
||||
"""Provider output could not be parsed or validated against the schema."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelMessage:
|
||||
role: Literal["system", "user", "assistant"]
|
||||
content: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelInvocationRequest:
|
||||
messages: Sequence[ModelMessage]
|
||||
model_role: str | None = None
|
||||
purpose: str | None = None
|
||||
response_schema: Mapping[str, object] | None = None
|
||||
timeout_seconds: float | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
object.__setattr__(self, "messages", tuple(self.messages))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelUsage:
|
||||
input_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
total_tokens: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelInvocationResult:
|
||||
content: str
|
||||
structured_output: Mapping[str, object] | None = None
|
||||
resolved_model: str | None = None
|
||||
usage: ModelUsage | None = None
|
||||
|
||||
|
||||
class ModelInvoker(Protocol):
|
||||
async def invoke(self, request: ModelInvocationRequest) -> ModelInvocationResult:
|
||||
"""Invoke once; schema requests succeed only with validated object data.
|
||||
|
||||
Caller cancellation propagates as asyncio.CancelledError. A timeout
|
||||
(including time waiting for host capacity) raises ModelInvocationFailed.
|
||||
"""
|
||||
raise ModelInvocationUnavailable("Model invocation is not supported by this host")
|
||||
@ -277,9 +277,34 @@ detached task store, the same fallback `notify_system_model_call` uses when its
|
||||
supplies none.
|
||||
|
||||
Gateway services start in registration order after the persistence engine and session
|
||||
factory are ready. Each receives the same `ExtensionRuntimeDeps` snapshot containing the
|
||||
factory are ready. Ungranted services share an `ExtensionRuntimeDeps` snapshot containing the
|
||||
app store, projected host policy, session factory, and optional read-only
|
||||
`RunEvidenceReader`. The Gateway constructs the configured run and event stores before
|
||||
`RunEvidenceReader`. `plugins[].host_access.model_invocation` optionally binds a model invoker
|
||||
to each service via the host-only `ModelInvocationService` adapter. The loader captures
|
||||
one `ModelInvocationScope` per installation, not per `use` string, so duplicate sources
|
||||
cannot inherit one another's roles. Its semaphore is shared by that installation's
|
||||
services; its admission ceiling is twice the concurrency limit, checked before
|
||||
payload processing. Provider work is shielded from caller cancellation and retains
|
||||
both budgets until actual completion, including synchronous LangChain executor calls
|
||||
and offloaded construction. Abandoned construction cannot dispatch a model request.
|
||||
Provider-task cancellation is a normalized failure; only a new cancellation of
|
||||
the invoking task propagates. Compare cancellation counts against invocation entry
|
||||
so previously handled caller cancellations do not mask provider failures.
|
||||
Failed-install positional rollback also removes its adapters. The adapter
|
||||
receives startup config through `start_with_host`, while extensions receive only the
|
||||
neutral invoker in a replaced deps snapshot. No-grant services preserve their old path.
|
||||
Failed start and stop revoke the service's handle and cancel queued/in-flight
|
||||
callers; they do not release slots owned by still-running provider work. Structured
|
||||
schema checks and output validation run in terminable isolated Python children,
|
||||
with pipe I/O on admission-bounded dedicated threads (Windows selector-loop compatible,
|
||||
independent of a potentially saturated provider executor). Cancellation kills and
|
||||
reaps those children before releasing admission.
|
||||
Grants and model profiles are startup snapshots; changing them requires restarting the
|
||||
Gateway. Calls use the normal model factory and attributed tracing, return plain text,
|
||||
usage counts and optionally locally validated JSON objects, and never return raw model
|
||||
objects or provider exception chains. See `backend/docs/extension-model-invocation.md`.
|
||||
|
||||
The Gateway constructs the configured run and event stores before
|
||||
services so the reader is usable from `start()`. Changed-run discovery uses an opaque,
|
||||
scope-bound cursor over `(change_seq, run_id)`; a run that changes after it was returned may
|
||||
be replayed, but an unreturned run cannot be skipped. Legacy rows start at `change_seq=0`
|
||||
|
||||
@ -22,6 +22,7 @@ from typing import Any
|
||||
from deerflow_extension_api import ExtensionRuntimeDeps
|
||||
|
||||
from deerflow.extensions.loader import Diagnostic
|
||||
from deerflow.extensions.model_access import ModelInvocationService
|
||||
from deerflow.extensions.policy import project_host_policy
|
||||
from deerflow.extensions.registry import LoadedExtensions
|
||||
|
||||
@ -599,7 +600,10 @@ async def start_services(
|
||||
attempted_services.append(entry)
|
||||
cancellation_count = _cancellation_count()
|
||||
try:
|
||||
await service.start(deps)
|
||||
if isinstance(service, ModelInvocationService):
|
||||
await service.start_with_host(deps, app_config)
|
||||
else:
|
||||
await service.start(deps)
|
||||
except asyncio.CancelledError:
|
||||
if _cancellation_count() > cancellation_count:
|
||||
raise
|
||||
|
||||
@ -16,6 +16,7 @@ from typing import Any, Literal
|
||||
from deerflow_extension_api import API_VERSION
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from deerflow.extensions.model_access import ExtensionHostAccess, ModelInvocationScope
|
||||
from deerflow.extensions.registry import ExtensionRegistry, LoadedExtensions
|
||||
from deerflow.persistence.migrations._env_filters import register_extension_table_prefix
|
||||
from deerflow.reflection import resolve_variable
|
||||
@ -43,6 +44,7 @@ class ExtensionSpec(BaseModel):
|
||||
description="Installed Python distribution recorded by the extension manager",
|
||||
)
|
||||
use: str = Field(description="Entry point path, e.g. 'my_extension:install'")
|
||||
host_access: ExtensionHostAccess = Field(default_factory=ExtensionHostAccess)
|
||||
config: dict[str, Any] = Field(
|
||||
default_factory=dict,
|
||||
description="Extension-private configuration, passed to install() verbatim",
|
||||
@ -237,7 +239,9 @@ def load_extensions(specs: Sequence[ExtensionSpec]) -> tuple[LoadedExtensions, l
|
||||
# installed instance that happens to share this spec's `use`.
|
||||
mark = registry.mark()
|
||||
try:
|
||||
with registry.attributed_to(spec.use):
|
||||
grant = spec.host_access.model_invocation
|
||||
scope = ModelInvocationScope(spec.use, grant) if grant is not None else None
|
||||
with registry.attributed_to(spec.use, model_access=scope):
|
||||
install(registry, _frozen_config(spec.config))
|
||||
except Exception as exc:
|
||||
registry.rollback_to(mark)
|
||||
|
||||
94
backend/packages/harness/deerflow/extensions/model_access.py
Normal file
94
backend/packages/harness/deerflow/extensions/model_access.py
Normal file
@ -0,0 +1,94 @@
|
||||
"""Operator grants and service lifetime binding; no provider imports at startup."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import replace
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
|
||||
class ModelInvocationGrant(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False)
|
||||
|
||||
roles: dict[str, str] = Field(min_length=1, description="Logical role to configured model name")
|
||||
max_concurrency: int = Field(default=2, ge=1, le=64)
|
||||
timeout_seconds: float = Field(default=60, gt=0, le=600)
|
||||
max_input_chars: int = Field(default=262144, ge=1, le=1048576)
|
||||
max_output_chars: int = Field(default=65536, ge=1, le=1048576)
|
||||
|
||||
@field_validator("roles")
|
||||
@classmethod
|
||||
def validate_roles(cls, roles):
|
||||
if any(not value or value != value.strip() for pair in roles.items() for value in pair):
|
||||
raise ValueError("Role and model names must be nonempty without surrounding whitespace")
|
||||
return roles
|
||||
|
||||
|
||||
class ExtensionHostAccess(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
model_invocation: ModelInvocationGrant | None = None
|
||||
|
||||
|
||||
class ModelInvocationScope:
|
||||
"""One install() owns one budget, even if it registers multiple services."""
|
||||
|
||||
def __init__(self, source: str, grant: ModelInvocationGrant):
|
||||
self.source = source
|
||||
self.grant = grant.model_copy(deep=True)
|
||||
self.budget: ModelInvocationBudget | None = None
|
||||
|
||||
def bind(self, app_config: Any):
|
||||
from deerflow.extensions.model_invocation import HostModelInvoker
|
||||
|
||||
if self.budget is None:
|
||||
self.budget = ModelInvocationBudget(self.grant.max_concurrency)
|
||||
return HostModelInvoker(self.source, self.grant, self.budget, app_config)
|
||||
|
||||
|
||||
class ModelInvocationBudget:
|
||||
"""Loop-owned admission shared by every service in one installation."""
|
||||
|
||||
def __init__(self, concurrency: int):
|
||||
self.semaphore = asyncio.Semaphore(concurrency)
|
||||
self.capacity = 2 * concurrency
|
||||
self.admitted = 0
|
||||
# Strong references keep abandoned, shielded provider work alive.
|
||||
self.workers: set[asyncio.Task] = set()
|
||||
|
||||
def retain(self, task):
|
||||
self.workers.add(task)
|
||||
task.add_done_callback(self._finished)
|
||||
|
||||
def _finished(self, task):
|
||||
self.workers.discard(task)
|
||||
if not task.cancelled():
|
||||
task.exception() # Consume abandoned failures without exposing provider text.
|
||||
|
||||
|
||||
class ModelInvocationService:
|
||||
"""Host adapter keeps capability revocation paired with service cleanup."""
|
||||
|
||||
def __init__(self, service: Any, scope: ModelInvocationScope):
|
||||
self.service = service
|
||||
self.scope = scope
|
||||
self.invoker = None
|
||||
|
||||
async def start(self, deps):
|
||||
raise RuntimeError("Granted services must be started through the host lifecycle")
|
||||
|
||||
async def start_with_host(self, deps, app_config):
|
||||
self.invoker = self.scope.bind(app_config)
|
||||
try:
|
||||
await self.service.start(replace(deps, model_invoker=self.invoker))
|
||||
except BaseException:
|
||||
# Includes host cancellation: start() may already have spawned calls.
|
||||
self.invoker.close()
|
||||
raise
|
||||
|
||||
async def stop(self):
|
||||
if self.invoker is not None:
|
||||
self.invoker.close()
|
||||
await self.service.stop()
|
||||
263
backend/packages/harness/deerflow/extensions/model_invocation.py
Normal file
263
backend/packages/harness/deerflow/extensions/model_invocation.py
Normal file
@ -0,0 +1,263 @@
|
||||
"""Host-owned model construction, bounded invocation and neutral projection."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from deerflow_extension_api import (
|
||||
ModelInvocationError,
|
||||
ModelInvocationFailed,
|
||||
ModelInvocationRequest,
|
||||
ModelInvocationResult,
|
||||
ModelInvocationUnauthorized,
|
||||
ModelInvocationUnavailable,
|
||||
ModelMessage,
|
||||
ModelOutputValidationError,
|
||||
ModelUsage,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
||||
|
||||
from deerflow.models import create_chat_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_MESSAGE_TYPES = {"system": SystemMessage, "user": HumanMessage, "assistant": AIMessage}
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Invocation:
|
||||
acquired: bool = False
|
||||
abandoned: bool = False
|
||||
provider: asyncio.Task | None = None
|
||||
|
||||
|
||||
async def _schema_validator(schema_json, content=None):
|
||||
"""Run CPU-bound JSON Schema work outside the Gateway process/GIL.
|
||||
|
||||
A dedicated pipe-I/O thread also works with Windows' selector loop and cannot
|
||||
be starved by sync providers occupying asyncio's default executor. Admission
|
||||
bounds both these threads and child processes. Cancellation kills and reaps
|
||||
the child before this call releases its slot.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
finished = loop.create_future()
|
||||
cancelled = threading.Event()
|
||||
payload = schema_json + "\n"
|
||||
if content is not None:
|
||||
payload += json.dumps(content) + "\n"
|
||||
|
||||
def work():
|
||||
status = b""
|
||||
try:
|
||||
options = {"creationflags": subprocess.CREATE_NO_WINDOW} if sys.platform == "win32" else {}
|
||||
with subprocess.Popen(
|
||||
[sys.executable, "-I", str(Path(__file__).with_name("model_schema_worker.py"))],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.DEVNULL,
|
||||
**options,
|
||||
) as process:
|
||||
try:
|
||||
data = payload.encode("utf-8")
|
||||
while not cancelled.is_set():
|
||||
try:
|
||||
status, _ = process.communicate(data, timeout=0.05)
|
||||
break
|
||||
except subprocess.TimeoutExpired:
|
||||
data = None # communicate retains pending input for the next poll.
|
||||
finally:
|
||||
if process.poll() is None:
|
||||
process.kill()
|
||||
process.communicate()
|
||||
except Exception:
|
||||
pass # Only fixed validation status crosses the public error boundary.
|
||||
finally:
|
||||
loop.call_soon_threadsafe(finished.set_result, status.split())
|
||||
|
||||
threading.Thread(target=work, name="extension-schema-validation", daemon=True).start()
|
||||
interrupted = False
|
||||
while not finished.done():
|
||||
try:
|
||||
await asyncio.shield(finished)
|
||||
except asyncio.CancelledError:
|
||||
interrupted = True
|
||||
cancelled.set()
|
||||
if interrupted:
|
||||
raise asyncio.CancelledError
|
||||
status = finished.result()
|
||||
if content is None:
|
||||
if status != [b"ready"]:
|
||||
raise ModelInvocationFailed("response_schema must be a valid inline Draft 2020-12 object schema")
|
||||
elif status != [b"ready", b"valid"]:
|
||||
raise ModelOutputValidationError("Model output did not match response_schema")
|
||||
|
||||
|
||||
def _usage(metadata):
|
||||
if not isinstance(metadata, Mapping):
|
||||
return None
|
||||
|
||||
def token_count(key):
|
||||
value = metadata.get(key)
|
||||
return value if type(value) is int and value >= 0 else None
|
||||
|
||||
return ModelUsage(token_count("input_tokens"), token_count("output_tokens"), token_count("total_tokens"))
|
||||
|
||||
|
||||
class HostModelInvoker:
|
||||
def __init__(self, source, grant, budget, app_config):
|
||||
self._source = source
|
||||
self._grant = grant
|
||||
self._budget = budget
|
||||
self._app_config = app_config
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._closed = False
|
||||
self._tasks = set()
|
||||
|
||||
def close(self):
|
||||
"""Revoke handles and cancel callers; running providers retain their slots."""
|
||||
self._closed = True
|
||||
for task in tuple(self._tasks):
|
||||
task.cancel()
|
||||
|
||||
async def invoke(self, request: ModelInvocationRequest) -> ModelInvocationResult:
|
||||
if self._closed:
|
||||
raise ModelInvocationUnavailable("Model invocation capability has stopped")
|
||||
if asyncio.get_running_loop() is not self._loop:
|
||||
raise ModelInvocationUnavailable("Model invocation requires the service event loop")
|
||||
if self._budget.admitted >= self._budget.capacity:
|
||||
raise ModelInvocationFailed("Model invocation capacity exceeded")
|
||||
self._budget.admitted += 1
|
||||
call = _Invocation()
|
||||
deadline = None
|
||||
task = asyncio.current_task()
|
||||
cancellation_count = task.cancelling()
|
||||
self._tasks.add(task)
|
||||
try:
|
||||
if not isinstance(request, ModelInvocationRequest):
|
||||
raise ModelInvocationFailed("Expected ModelInvocationRequest")
|
||||
timeout = request.timeout_seconds
|
||||
if timeout is not None and (type(timeout) not in (int, float) or not math.isfinite(timeout) or timeout <= 0):
|
||||
raise ModelInvocationFailed("timeout_seconds must be finite and positive")
|
||||
timeout = min(timeout, self._grant.timeout_seconds) if timeout is not None else self._grant.timeout_seconds
|
||||
deadline = asyncio.timeout(timeout)
|
||||
async with deadline:
|
||||
await self._budget.semaphore.acquire()
|
||||
call.acquired = True
|
||||
return await self._invoke(request, call)
|
||||
except asyncio.CancelledError:
|
||||
# shield also raises when only the provider task is cancelled. A
|
||||
# previously handled caller cancellation must not mask that failure.
|
||||
if task.cancelling() > cancellation_count or call.provider is None or not call.provider.cancelled():
|
||||
raise
|
||||
failure = ModelInvocationFailed("Model provider cancelled")
|
||||
except ModelInvocationError as exc:
|
||||
# Copy only our normalized text. In particular, do not pass a JSON,
|
||||
# schema or provider exception through __context__ to extensions.
|
||||
failure = type(exc)(str(exc))
|
||||
except TimeoutError:
|
||||
failure = ModelInvocationFailed("Model invocation timed out" if deadline is not None and deadline.expired() else "Model provider timed out")
|
||||
except Exception as exc:
|
||||
logger.warning("Extension model invocation failed for %s (%s)", self._source, type(exc).__name__)
|
||||
failure = ModelInvocationFailed("Model invocation failed")
|
||||
finally:
|
||||
self._tasks.discard(task)
|
||||
call.abandoned = True
|
||||
|
||||
def release(_=None):
|
||||
self._budget.admitted -= 1
|
||||
if call.acquired:
|
||||
self._budget.semaphore.release()
|
||||
|
||||
if call.provider is not None and not call.provider.done():
|
||||
call.provider.add_done_callback(release)
|
||||
else:
|
||||
release()
|
||||
raise failure from None
|
||||
|
||||
async def _provider_call(self, name, messages, role, purpose, call):
|
||||
# Neither to_thread construction nor a sync-backed ainvoke is cancellable.
|
||||
# The caller may leave, but the real work keeps its admission/slot until done.
|
||||
model = await asyncio.to_thread(create_chat_model, name, app_config=self._app_config)
|
||||
if self._closed or call.abandoned:
|
||||
raise ModelInvocationUnavailable("Model invocation capability has stopped")
|
||||
return await model.ainvoke(
|
||||
messages,
|
||||
config={"run_name": "extension_model_invocation", "metadata": {"extension_source": self._source, "extension_model_role": role, "extension_purpose": purpose}},
|
||||
)
|
||||
|
||||
async def _invoke(self, request, call):
|
||||
role = request.model_role if request.model_role is not None else "default"
|
||||
if not isinstance(role, str) or role not in self._grant.roles:
|
||||
raise ModelInvocationUnauthorized("Model role is not granted to this extension")
|
||||
name = self._grant.roles[role]
|
||||
if self._app_config.get_model_config(name) is None:
|
||||
raise ModelInvocationUnavailable("Configured model is unavailable")
|
||||
if request.purpose is not None and (not isinstance(request.purpose, str) or not 1 <= len(request.purpose) <= 128):
|
||||
raise ModelInvocationFailed("purpose must contain 1 to 128 characters")
|
||||
if not request.messages or len(request.messages) > 256:
|
||||
raise ModelInvocationFailed("Expected 1 to 256 text messages")
|
||||
messages = []
|
||||
input_chars = 0
|
||||
for message in request.messages:
|
||||
if not isinstance(message, ModelMessage) or not isinstance(message.role, str) or message.role not in _MESSAGE_TYPES or not isinstance(message.content, str):
|
||||
raise ModelInvocationFailed("Only system, user and assistant text messages are supported")
|
||||
input_chars += len(message.content)
|
||||
if input_chars > self._grant.max_input_chars:
|
||||
raise ModelInvocationFailed("Model input exceeds host limit")
|
||||
messages.append(_MESSAGE_TYPES[message.role](content=message.content))
|
||||
|
||||
schema_json = None
|
||||
if request.response_schema is not None:
|
||||
if not isinstance(request.response_schema, Mapping):
|
||||
raise ModelInvocationFailed("response_schema must be an object schema")
|
||||
# Snapshot the schema before validation crosses the process boundary.
|
||||
schema_json = json.dumps(dict(request.response_schema), allow_nan=False)
|
||||
instruction = "Return only a JSON object matching the supplied response schema (no Markdown)."
|
||||
schema_data = "Required response schema (JSON data):\n" + schema_json
|
||||
if input_chars + len(instruction) + len(schema_data) > self._grant.max_input_chars:
|
||||
raise ModelInvocationFailed("Model input including schema exceeds host limit")
|
||||
messages.insert(0, SystemMessage(content=instruction))
|
||||
# Schema descriptions may contain extension-owned input. Keep them
|
||||
# out of the host's fixed system instruction.
|
||||
messages.append(HumanMessage(content=schema_data))
|
||||
await _schema_validator(schema_json)
|
||||
|
||||
if self._closed:
|
||||
raise ModelInvocationUnavailable("Model invocation capability has stopped")
|
||||
call.provider = asyncio.create_task(self._provider_call(name, messages, role, request.purpose, call))
|
||||
self._budget.retain(call.provider)
|
||||
response = await asyncio.shield(call.provider)
|
||||
|
||||
if self._closed:
|
||||
raise ModelInvocationUnavailable("Model invocation capability has stopped")
|
||||
|
||||
if getattr(response, "tool_calls", None) or getattr(response, "invalid_tool_calls", None):
|
||||
raise ModelInvocationFailed("Tool-call responses are not supported")
|
||||
content = response.content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for block in content:
|
||||
if isinstance(block, str):
|
||||
parts.append(block)
|
||||
elif isinstance(block, dict) and block.get("type") == "text" and isinstance(block.get("text"), str):
|
||||
parts.append(block["text"])
|
||||
else:
|
||||
raise ModelInvocationFailed("Model returned non-text content")
|
||||
content = "".join(parts)
|
||||
if not isinstance(content, str):
|
||||
raise ModelInvocationFailed("Model returned non-text content")
|
||||
if len(content) > self._grant.max_output_chars:
|
||||
raise ModelInvocationFailed("Model output exceeds host limit")
|
||||
structured = None
|
||||
if schema_json is not None:
|
||||
await _schema_validator(schema_json, content)
|
||||
structured = json.loads(content)
|
||||
return ModelInvocationResult(content, structured, name, _usage(getattr(response, "usage_metadata", None)))
|
||||
@ -0,0 +1,57 @@
|
||||
"""One-call schema worker, launched by file path in an isolated Python process.
|
||||
|
||||
Only bounded JSON input and fixed status words cross the pipes. Expensive schema
|
||||
checking and regex validation must never execute on the Gateway loop or its GIL.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
|
||||
from jsonschema import Draft202012Validator
|
||||
|
||||
|
||||
def _reject_constant(value):
|
||||
raise ValueError("Non-finite JSON number")
|
||||
|
||||
|
||||
def _check_inline(value):
|
||||
if isinstance(value, dict):
|
||||
if any(key in value for key in ("$ref", "$dynamicRef", "$recursiveRef")):
|
||||
raise ValueError("References are unsupported")
|
||||
if "$schema" in value and value["$schema"] != "https://json-schema.org/draft/2020-12/schema":
|
||||
raise ValueError("Unsupported dialect")
|
||||
for child in value.values():
|
||||
_check_inline(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
_check_inline(child)
|
||||
|
||||
|
||||
def main():
|
||||
try:
|
||||
schema = json.loads(sys.stdin.buffer.readline())
|
||||
if schema.get("type") != "object":
|
||||
raise ValueError("Expected object schema")
|
||||
_check_inline(schema)
|
||||
Draft202012Validator.check_schema(schema)
|
||||
validator = Draft202012Validator(schema)
|
||||
except Exception:
|
||||
print("invalid_schema", flush=True)
|
||||
return
|
||||
print("ready", flush=True)
|
||||
response = sys.stdin.buffer.readline()
|
||||
if not response:
|
||||
return
|
||||
try:
|
||||
content = json.loads(response)
|
||||
structured = json.loads(content, parse_constant=_reject_constant)
|
||||
json.dumps(structured, allow_nan=False)
|
||||
validator.validate(structured)
|
||||
except Exception:
|
||||
print("invalid_output", flush=True)
|
||||
return
|
||||
print("valid", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -26,6 +26,8 @@ from deerflow_extension_api import (
|
||||
from deerflow_extension_api import ExtensionRegistry as ExtensionRegistryContract
|
||||
from deerflow_extension_api.plugins import BrowserAssets, BrowserModule, PluginContribution
|
||||
|
||||
from deerflow.extensions.model_access import ModelInvocationScope, ModelInvocationService
|
||||
|
||||
_Entry = tuple[str, Any]
|
||||
|
||||
|
||||
@ -83,16 +85,20 @@ class ExtensionRegistry(ExtensionRegistryContract):
|
||||
self._routers: list[_Entry] = []
|
||||
self._plugins: list[_Entry] = []
|
||||
self._current_source: str | None = None
|
||||
self._model_access: ModelInvocationScope | None = None
|
||||
|
||||
@contextmanager
|
||||
def attributed_to(self, source: str) -> Iterator[None]:
|
||||
def attributed_to(self, source: str, *, model_access: ModelInvocationScope | None = None) -> Iterator[None]:
|
||||
"""Attribute everything registered inside the block to ``source``."""
|
||||
previous = self._current_source
|
||||
previous_access = self._model_access
|
||||
self._current_source = source
|
||||
self._model_access = model_access
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self._current_source = previous
|
||||
self._model_access = previous_access
|
||||
|
||||
def _source(self) -> str:
|
||||
if self._current_source is None:
|
||||
@ -154,6 +160,8 @@ class ExtensionRegistry(ExtensionRegistryContract):
|
||||
self._context_compaction_observers.append((self._source(), observer))
|
||||
|
||||
def service(self, service: ExtensionService) -> None:
|
||||
if self._model_access is not None:
|
||||
service = ModelInvocationService(service, self._model_access)
|
||||
self._services.append((self._source(), service))
|
||||
|
||||
def routers(self, routers: Sequence[Any]) -> None:
|
||||
|
||||
57
backend/tests/blocking_io/test_extension_model_invocation.py
Normal file
57
backend/tests/blocking_io/test_extension_model_invocation.py
Normal file
@ -0,0 +1,57 @@
|
||||
"""Provider construction may read credentials but must never block the host loop."""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from deerflow_extension_api import ModelInvocationRequest, ModelMessage
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
from deerflow.extensions import model_invocation
|
||||
from deerflow.extensions.model_access import ModelInvocationBudget, ModelInvocationGrant
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_construction_runs_off_loop(tmp_path, monkeypatch):
|
||||
credentials = tmp_path / "provider.txt"
|
||||
await asyncio.to_thread(credentials.write_text, "test-only", encoding="utf-8")
|
||||
threads = []
|
||||
|
||||
def factory(name, *, app_config):
|
||||
assert credentials.read_text(encoding="utf-8") == "test-only"
|
||||
threads.append(threading.get_ident())
|
||||
return SimpleNamespace(ainvoke=AsyncMock(return_value=AIMessage(content="ok")))
|
||||
|
||||
monkeypatch.setattr(model_invocation, "create_chat_model", factory)
|
||||
invoker = model_invocation.HostModelInvoker(
|
||||
"example:install",
|
||||
ModelInvocationGrant(roles={"default": "model"}),
|
||||
ModelInvocationBudget(1),
|
||||
SimpleNamespace(get_model_config=lambda _: object()),
|
||||
)
|
||||
result = await invoker.invoke(ModelInvocationRequest([ModelMessage("user", "hello")]))
|
||||
assert result.content == "ok"
|
||||
assert threads and threads[0] != threading.get_ident()
|
||||
invoker.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schema_subprocess_io_runs_off_loop(monkeypatch):
|
||||
model = SimpleNamespace(ainvoke=AsyncMock(return_value=AIMessage(content='{"label":"你好"}')))
|
||||
monkeypatch.setattr(model_invocation, "create_chat_model", lambda *args, **kwargs: model)
|
||||
invoker = model_invocation.HostModelInvoker(
|
||||
"example:install",
|
||||
ModelInvocationGrant(roles={"default": "model"}),
|
||||
ModelInvocationBudget(1),
|
||||
SimpleNamespace(get_model_config=lambda _: object()),
|
||||
)
|
||||
result = await invoker.invoke(
|
||||
ModelInvocationRequest(
|
||||
[ModelMessage("user", "hello")],
|
||||
response_schema={"type": "object", "properties": {"label": {"type": "string"}}},
|
||||
)
|
||||
)
|
||||
assert result.structured_output == {"label": "你好"}
|
||||
invoker.close()
|
||||
590
backend/tests/test_extension_model_invocation.py
Normal file
590
backend/tests/test_extension_model_invocation.py
Normal file
@ -0,0 +1,590 @@
|
||||
"""Granted extension calls exercise the loader, service lifecycle and host adapter."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from deerflow_extension_api import (
|
||||
ExtensionRuntimeDeps,
|
||||
ModelInvocationFailed,
|
||||
ModelInvocationRequest,
|
||||
ModelInvocationUnauthorized,
|
||||
ModelInvocationUnavailable,
|
||||
ModelMessage,
|
||||
ModelOutputValidationError,
|
||||
)
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
from deerflow.extensions.gateway import start_services, stop_services
|
||||
from deerflow.extensions.loader import ExtensionSpec, load_extensions
|
||||
|
||||
|
||||
class Service:
|
||||
async def start(self, deps):
|
||||
self.deps = deps
|
||||
|
||||
async def stop(self):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def host(monkeypatch):
|
||||
from deerflow.extensions import model_invocation
|
||||
|
||||
model = SimpleNamespace(ainvoke=AsyncMock(return_value=AIMessage(content='{"label":"positive"}', usage_metadata={"input_tokens": 8, "output_tokens": 4, "total_tokens": 12})))
|
||||
factory = Mock(return_value=model)
|
||||
monkeypatch.setattr(model_invocation, "create_chat_model", factory)
|
||||
config = SimpleNamespace(get_model_config=lambda name: object() if name == "host-model" else None)
|
||||
|
||||
async def start(grants, *, services_per_install=1, service_type=Service):
|
||||
services = []
|
||||
|
||||
def install(registry, config):
|
||||
for _ in range(services_per_install):
|
||||
service = service_type()
|
||||
services.append(service)
|
||||
registry.service(service)
|
||||
|
||||
monkeypatch.setattr("deerflow.extensions.loader.resolve_variable", lambda _: install)
|
||||
specs = [ExtensionSpec(use="example:install", host_access={"model_invocation": grant} if grant else {}) for grant in grants]
|
||||
loaded, diagnostics = load_extensions(specs)
|
||||
assert not diagnostics
|
||||
diagnostics = await start_services(loaded, config, None)
|
||||
return loaded, services, diagnostics
|
||||
|
||||
return SimpleNamespace(start=start, model=model, factory=factory)
|
||||
|
||||
|
||||
GRANT = {"roles": {"default": "host-model"}}
|
||||
SCHEMA = {"type": "object", "properties": {"label": {"enum": ["positive", "negative"]}}, "required": ["label"]}
|
||||
|
||||
|
||||
def request(**kwargs):
|
||||
return ModelInvocationRequest(messages=[ModelMessage("user", "Classify this text")], **kwargs)
|
||||
|
||||
|
||||
def test_old_extension_has_no_capability():
|
||||
assert ExtensionRuntimeDeps().model_invoker is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grant_is_bound_to_installation_not_entrypoint(host):
|
||||
loaded, services, diagnostics = await host.start([GRANT, None, {"roles": {"fast": "host-model"}}])
|
||||
assert not diagnostics
|
||||
assert services[1].deps.model_invoker is None
|
||||
result = await services[0].deps.model_invoker.invoke(request(response_schema=SCHEMA, purpose="classify"))
|
||||
assert result.content == '{"label":"positive"}'
|
||||
assert result.structured_output == {"label": "positive"}
|
||||
assert result.resolved_model == "host-model"
|
||||
assert result.usage.total_tokens == 12
|
||||
with pytest.raises(ModelInvocationUnauthorized):
|
||||
await services[2].deps.model_invoker.invoke(request())
|
||||
assert host.factory.call_count == 1
|
||||
metadata = host.model.ainvoke.call_args.kwargs["config"]["metadata"]
|
||||
assert metadata["extension_source"] == "example:install"
|
||||
assert metadata["extension_purpose"] == "classify"
|
||||
messages = host.model.ainvoke.call_args.args[0]
|
||||
assert messages[0].type == "system"
|
||||
assert '"properties"' not in messages[0].content
|
||||
assert messages[-1].type == "human"
|
||||
assert '"properties"' in messages[-1].content
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("content", ["not JSON", '{"label":"unknown"}', "[]", '{"label":"positive", "score":NaN}', '{"label":"positive", "score":1e999}'])
|
||||
async def test_schema_failure_never_returns_success(host, content):
|
||||
host.model.ainvoke.return_value = AIMessage(content=content)
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
with pytest.raises(ModelOutputValidationError):
|
||||
await services[0].deps.model_invoker.invoke(request(response_schema=SCHEMA))
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_errors_do_not_expose_credentials(host):
|
||||
host.model.ainvoke.side_effect = RuntimeError("secret-key in provider URL")
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
with pytest.raises(ModelInvocationFailed) as error:
|
||||
await services[0].deps.model_invoker.invoke(request())
|
||||
assert "secret-key" not in str(error.value)
|
||||
assert error.value.__cause__ is None
|
||||
assert error.value.__context__ is None
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_model_and_disallowed_role_fail_before_provider(host):
|
||||
loaded, services, _ = await host.start([{"roles": {"default": "missing"}}])
|
||||
with pytest.raises(ModelInvocationUnavailable):
|
||||
await services[0].deps.model_invoker.invoke(request())
|
||||
with pytest.raises(ModelInvocationUnauthorized):
|
||||
await services[0].deps.model_invoker.invoke(request(model_role="host-model"))
|
||||
host.factory.assert_not_called()
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrency_shared_across_services_and_queue_timeout(host):
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def invoke(*args, **kwargs):
|
||||
started.set()
|
||||
await release.wait()
|
||||
return AIMessage(content="ok")
|
||||
|
||||
host.model.ainvoke.side_effect = invoke
|
||||
loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}], services_per_install=2)
|
||||
first = asyncio.create_task(services[0].deps.model_invoker.invoke(request()))
|
||||
await started.wait()
|
||||
with pytest.raises(ModelInvocationFailed, match="timed out"):
|
||||
await services[1].deps.model_invoker.invoke(request(timeout_seconds=0.01))
|
||||
assert host.model.ainvoke.call_count == 1
|
||||
first.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await first
|
||||
release.set()
|
||||
assert (await services[1].deps.model_invoker.invoke(request())).content == "ok"
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_revokes_retained_capability_and_cancels_inflight(host):
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def invoke(*args, **kwargs):
|
||||
started.set()
|
||||
await release.wait()
|
||||
|
||||
host.model.ainvoke.side_effect = invoke
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
invoker = services[0].deps.model_invoker
|
||||
task = asyncio.create_task(invoker.invoke(request()))
|
||||
await started.wait()
|
||||
await stop_services(loaded)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
with pytest.raises(ModelInvocationUnavailable):
|
||||
await invoker.invoke(request())
|
||||
release.set()
|
||||
await asyncio.gather(*invoker._budget.workers)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_start_failure_revokes_only_failed_service(host):
|
||||
class Broken(Service):
|
||||
async def start(self, deps):
|
||||
await super().start(deps)
|
||||
raise ValueError("broken")
|
||||
|
||||
instances = iter([Broken(), Service()])
|
||||
loaded, services, diagnostics = await host.start([GRANT], services_per_install=2, service_type=instances.__next__)
|
||||
assert len(diagnostics) == 1
|
||||
with pytest.raises(ModelInvocationUnavailable):
|
||||
await services[0].deps.model_invoker.invoke(request())
|
||||
assert (await services[1].deps.model_invoker.invoke(request())).content == '{"label":"positive"}'
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_classification_example_uses_public_contract(host):
|
||||
from pathlib import Path
|
||||
|
||||
doc = (Path(__file__).resolve().parents[1] / "docs" / "extension-model-invocation.md").read_text(encoding="utf-8")
|
||||
code = doc.split("```python\n", 1)[1].split("```", 1)[0]
|
||||
namespace = {}
|
||||
exec(compile(code, "extension-model-invocation.md", "exec"), namespace)
|
||||
loaded, services, _ = await host.start([GRANT], service_type=namespace["Classifier"])
|
||||
assert await services[0].classify("Excellent work") == "positive"
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_host_factory_preserves_tracing_and_text_contract(monkeypatch):
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.config.model_config import ModelConfig
|
||||
from deerflow.config.sandbox_config import SandboxConfig
|
||||
|
||||
traces = []
|
||||
|
||||
class Observer(BaseCallbackHandler):
|
||||
def on_chat_model_start(self, serialized, messages, **kwargs):
|
||||
traces.append(kwargs["metadata"])
|
||||
|
||||
config = AppConfig(
|
||||
models=[ModelConfig(name="host-model", model="fake", use="langchain_core.language_models.fake_chat_models:FakeListChatModel", responses=['{"label":"negative"}'])],
|
||||
sandbox=SandboxConfig(use="deerflow.sandbox.local:LocalSandboxProvider"),
|
||||
)
|
||||
service = Service()
|
||||
monkeypatch.setattr("deerflow.extensions.loader.resolve_variable", lambda _: lambda registry, _: registry.service(service))
|
||||
monkeypatch.setattr("deerflow.models.factory.build_tracing_callbacks", lambda: [Observer()])
|
||||
loaded, diagnostics = load_extensions([ExtensionSpec(use="real:install", host_access={"model_invocation": GRANT})])
|
||||
assert not diagnostics
|
||||
assert not await start_services(loaded, config, None)
|
||||
result = await service.deps.model_invoker.invoke(request(response_schema=SCHEMA))
|
||||
assert result.structured_output == {"label": "negative"}
|
||||
assert traces[0]["extension_source"] == "real:install"
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"schema",
|
||||
[
|
||||
{"type": "object", "properties": {"x": {"$ref": "https://example.com/schema"}}},
|
||||
{"type": "object", "$ref": "#/properties/x"},
|
||||
{"type": "array"},
|
||||
{"type": "object", "required": "label"},
|
||||
{"type": "object", "$schema": "http://json-schema.org/draft-07/schema#"},
|
||||
],
|
||||
)
|
||||
async def test_invalid_schema_rejected_without_provider_call(host, schema):
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
with pytest.raises(ModelInvocationFailed):
|
||||
await services[0].deps.model_invoker.invoke(request(response_schema=schema))
|
||||
host.factory.assert_not_called()
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf"), True, "60"])
|
||||
async def test_invalid_timeout_rejected_before_provider(host, timeout):
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
with pytest.raises(ModelInvocationFailed):
|
||||
await services[0].deps.model_invoker.invoke(request(timeout_seconds=timeout))
|
||||
host.factory.assert_not_called()
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_timeout_caps_request_and_preserves_running_provider(host):
|
||||
cancelled = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def invoke(*args, **kwargs):
|
||||
try:
|
||||
await release.wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
host.model.ainvoke.side_effect = invoke
|
||||
loaded, services, _ = await host.start([{**GRANT, "timeout_seconds": 0.05}])
|
||||
with pytest.raises(ModelInvocationFailed, match="timed out"):
|
||||
await services[0].deps.model_invoker.invoke(request(timeout_seconds=500))
|
||||
assert not cancelled.is_set()
|
||||
release.set()
|
||||
await asyncio.wait_for(cancelled.wait(), 1)
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_and_output_limits(host):
|
||||
loaded, services, _ = await host.start([{**GRANT, "max_input_chars": 5, "max_output_chars": 2}])
|
||||
invoker = services[0].deps.model_invoker
|
||||
with pytest.raises(ModelInvocationFailed, match="input"):
|
||||
await invoker.invoke(request())
|
||||
host.factory.assert_not_called()
|
||||
with pytest.raises(ModelInvocationFailed, match="output"):
|
||||
await invoker.invoke(ModelInvocationRequest([ModelMessage("user", "hi")]))
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_text_blocks_project_without_provider_metadata(host):
|
||||
host.model.ainvoke.return_value = AIMessage(content=[{"type": "text", "text": "hello"}, {"type": "text", "text": " world"}], response_metadata={"secret": "hidden"})
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
result = await services[0].deps.model_invoker.invoke(request())
|
||||
assert result.content == "hello world"
|
||||
assert result.usage is None
|
||||
assert "hidden" not in repr(result)
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_calls_rejected(host):
|
||||
host.model.ainvoke.return_value = AIMessage(content="", tool_calls=[{"name": "search", "args": {}, "id": "call"}])
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
with pytest.raises(ModelInvocationFailed, match="Tool-call"):
|
||||
await services[0].deps.model_invoker.invoke(request())
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_duplicate_install_does_not_grant_prior_service(monkeypatch):
|
||||
good = Service()
|
||||
failed = Service()
|
||||
|
||||
def install(registry, config):
|
||||
registry.service(failed if config else good)
|
||||
if config:
|
||||
raise RuntimeError("install failed after registration")
|
||||
|
||||
monkeypatch.setattr("deerflow.extensions.loader.resolve_variable", lambda _: install)
|
||||
loaded, diagnostics = load_extensions(
|
||||
[
|
||||
ExtensionSpec(use="same:install"),
|
||||
ExtensionSpec(use="same:install", config={"fail": True}, host_access={"model_invocation": GRANT}),
|
||||
]
|
||||
)
|
||||
assert len(diagnostics) == 1
|
||||
assert len(loaded.services) == 1
|
||||
await start_services(loaded, SimpleNamespace(), None)
|
||||
assert good.deps.model_invoker is None
|
||||
assert not hasattr(failed, "deps")
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("grant", [{"roles": {}}, {"roles": {"default": " "}}, {"roles": {"": "model"}}, {**GRANT, "max_concurrency": 0}, {**GRANT, "timeout_seconds": float("inf")}])
|
||||
def test_invalid_grants_rejected(grant):
|
||||
with pytest.raises(ValueError):
|
||||
ExtensionSpec(use="example:install", host_access={"model_invocation": grant})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("exit_kind", ["timeout", "cancel", "stop"])
|
||||
async def test_sync_provider_retains_slot_until_thread_finishes(host, exit_kind):
|
||||
import threading
|
||||
|
||||
from langchain_core.language_models.chat_models import BaseChatModel
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
|
||||
started = asyncio.Event()
|
||||
release = threading.Event()
|
||||
loop = asyncio.get_running_loop()
|
||||
calls = []
|
||||
|
||||
class SyncModel(BaseChatModel):
|
||||
@property
|
||||
def _llm_type(self):
|
||||
return "blocked-sync-test"
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
calls.append(1)
|
||||
loop.call_soon_threadsafe(started.set)
|
||||
assert release.wait(10)
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content="ok"))])
|
||||
|
||||
host.factory.return_value = SyncModel()
|
||||
loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}], services_per_install=2)
|
||||
invoker = services[0].deps.model_invoker
|
||||
first = asyncio.create_task(invoker.invoke(request(timeout_seconds=0.1 if exit_kind == "timeout" else 5)))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 3)
|
||||
if exit_kind == "cancel":
|
||||
first.cancel()
|
||||
elif exit_kind == "stop":
|
||||
invoker.close()
|
||||
with pytest.raises(ModelInvocationFailed if exit_kind == "timeout" else asyncio.CancelledError):
|
||||
await first
|
||||
with pytest.raises(ModelInvocationFailed, match="timed out"):
|
||||
await services[1].deps.model_invoker.invoke(request(timeout_seconds=0.05))
|
||||
assert len(calls) == 1, "timed-out synchronous requests still consume their concurrency slot"
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(first, return_exceptions=True)
|
||||
assert (await services[1].deps.model_invoker.invoke(request())).content == "ok"
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_admission_limit_shared_and_rejects_before_payload_processing(host):
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def invoke(*args, **kwargs):
|
||||
started.set()
|
||||
await release.wait()
|
||||
return AIMessage(content="ok")
|
||||
|
||||
host.model.ainvoke.side_effect = invoke
|
||||
loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}], services_per_install=2)
|
||||
first = asyncio.create_task(services[0].deps.model_invoker.invoke(request()))
|
||||
await started.wait()
|
||||
queued = asyncio.create_task(services[1].deps.model_invoker.invoke(request()))
|
||||
await asyncio.sleep(0)
|
||||
try:
|
||||
with pytest.raises(ModelInvocationFailed, match="capacity"):
|
||||
await services[0].deps.model_invoker.invoke(request(timeout_seconds=0.05))
|
||||
assert host.factory.call_count == 1
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(first, queued)
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_provider_timeout_is_not_reported_as_host_deadline(host):
|
||||
host.model.ainvoke.side_effect = TimeoutError("secret provider URL")
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
with pytest.raises(ModelInvocationFailed, match="provider timed out") as error:
|
||||
await services[0].deps.model_invoker.invoke(request())
|
||||
assert error.value.__context__ is None
|
||||
assert "secret" not in str(error.value)
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("self_cancel", [False, True])
|
||||
@pytest.mark.parametrize("prior_cancellation", [False, True])
|
||||
async def test_provider_cancellation_is_normalized_and_releases_capacity(host, self_cancel, prior_cancellation):
|
||||
async def cancelled_provider(*args, **kwargs):
|
||||
if self_cancel:
|
||||
asyncio.current_task().cancel("secret provider detail")
|
||||
await asyncio.sleep(0)
|
||||
raise asyncio.CancelledError("secret provider detail")
|
||||
|
||||
host.model.ainvoke.side_effect = cancelled_provider
|
||||
loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}])
|
||||
invoker = services[0].deps.model_invoker
|
||||
|
||||
async def caller():
|
||||
if prior_cancellation:
|
||||
asyncio.current_task().cancel()
|
||||
try:
|
||||
await asyncio.sleep(0)
|
||||
except asyncio.CancelledError:
|
||||
pass # A previously handled request must not mask provider failure.
|
||||
with pytest.raises(ModelInvocationFailed, match="provider cancelled") as error:
|
||||
await invoker.invoke(request())
|
||||
assert error.value.__cause__ is None
|
||||
assert error.value.__context__ is None
|
||||
assert "secret" not in str(error.value)
|
||||
assert invoker._budget.admitted == 0
|
||||
host.model.ainvoke.side_effect = None
|
||||
assert (await invoker.invoke(request())).content == '{"label":"positive"}'
|
||||
|
||||
try:
|
||||
await asyncio.create_task(caller())
|
||||
finally:
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_cancellation_wins_when_provider_also_cancels(host):
|
||||
async def cancelled_provider(*args, **kwargs):
|
||||
caller.cancel()
|
||||
raise asyncio.CancelledError("provider cancelled too")
|
||||
|
||||
host.model.ainvoke.side_effect = cancelled_provider
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
invoker = services[0].deps.model_invoker
|
||||
caller = asyncio.create_task(invoker.invoke(request()))
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await caller
|
||||
assert caller.cancelled()
|
||||
assert invoker._budget.admitted == 0
|
||||
finally:
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_caller_cancellation_at_invocation_entry_propagates(host):
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
invoker = services[0].deps.model_invoker
|
||||
|
||||
async def caller():
|
||||
asyncio.current_task().cancel()
|
||||
await invoker.invoke(request())
|
||||
|
||||
try:
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.create_task(caller())
|
||||
finally:
|
||||
await asyncio.gather(*invoker._budget.workers, return_exceptions=True)
|
||||
await stop_services(loaded)
|
||||
assert invoker._budget.admitted == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("exit_kind", ["timeout", "cancel", "stop"])
|
||||
async def test_validation_deadline_does_not_block_event_loop(host, monkeypatch, exit_kind):
|
||||
import time
|
||||
|
||||
from deerflow.extensions import model_invocation
|
||||
|
||||
processes = []
|
||||
popen = model_invocation.subprocess.Popen
|
||||
|
||||
def track_process(*args, **kwargs):
|
||||
process = popen(*args, **kwargs)
|
||||
processes.append(process)
|
||||
return process
|
||||
|
||||
monkeypatch.setattr(model_invocation.subprocess, "Popen", track_process)
|
||||
schema = {"type": "object", "properties": {"value": {"type": "string", "pattern": "^(a+)+$"}}}
|
||||
host.model.ainvoke.return_value = AIMessage(content='{"value":"' + "a" * 30 + '!"}')
|
||||
loaded, services, _ = await host.start([GRANT])
|
||||
ticks = []
|
||||
running = True
|
||||
|
||||
async def heartbeat():
|
||||
while running:
|
||||
ticks.append(time.monotonic())
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
ticker = asyncio.create_task(heartbeat())
|
||||
invoker = services[0].deps.model_invoker
|
||||
task = asyncio.create_task(invoker.invoke(request(response_schema=schema, timeout_seconds=3)))
|
||||
try:
|
||||
if exit_kind != "timeout":
|
||||
async with asyncio.timeout(3):
|
||||
while len(processes) < 2:
|
||||
await asyncio.sleep(0.01)
|
||||
if exit_kind == "cancel":
|
||||
task.cancel()
|
||||
# Repeated cancellation must not abandon the child cleanup.
|
||||
asyncio.get_running_loop().call_later(0.01, task.cancel)
|
||||
else:
|
||||
invoker.close()
|
||||
with pytest.raises(ModelInvocationFailed if exit_kind == "timeout" else asyncio.CancelledError):
|
||||
await task
|
||||
assert host.model.ainvoke.call_count == 1, "deadline must fire during response validation"
|
||||
ticks.append(time.monotonic())
|
||||
assert max(b - a for a, b in zip(ticks, ticks[1:])) < 0.3
|
||||
assert len(processes) == 2
|
||||
assert all(process.poll() is not None for process in processes)
|
||||
assert invoker._budget.admitted == 0
|
||||
finally:
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
running = False
|
||||
await ticker
|
||||
await stop_services(loaded)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_constructor_retains_slot_and_never_dispatches(host):
|
||||
import threading
|
||||
|
||||
started = asyncio.Event()
|
||||
release = threading.Event()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def factory(*args, **kwargs):
|
||||
loop.call_soon_threadsafe(started.set)
|
||||
assert release.wait(10)
|
||||
return host.model
|
||||
|
||||
host.factory.side_effect = factory
|
||||
loaded, services, _ = await host.start([{**GRANT, "max_concurrency": 1}])
|
||||
invoker = services[0].deps.model_invoker
|
||||
first = asyncio.create_task(invoker.invoke(request()))
|
||||
try:
|
||||
await asyncio.wait_for(started.wait(), 3)
|
||||
first.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await first
|
||||
with pytest.raises(ModelInvocationFailed, match="timed out"):
|
||||
await invoker.invoke(request(timeout_seconds=0.05))
|
||||
assert host.factory.call_count == 1
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.gather(*invoker._budget.workers, return_exceptions=True)
|
||||
await stop_services(loaded)
|
||||
host.model.ainvoke.assert_not_called()
|
||||
@ -3222,6 +3222,14 @@ authorization:
|
||||
# enabled: true # false skips import and registration
|
||||
# required: false # true makes load failure abort Gateway startup
|
||||
# # (install --required opts in)
|
||||
# host_access: # optional host-owned grants (restart required)
|
||||
# model_invocation:
|
||||
# roles: # logical role -> configured models[].name
|
||||
# default: my-configured-model
|
||||
# max_concurrency: 2 # shared by this installation's services
|
||||
# timeout_seconds: 60 # includes waiting for capacity
|
||||
# max_input_chars: 262144 # includes the response-schema instruction
|
||||
# max_output_chars: 65536 # returned text, checked after generation
|
||||
# table_prefix: example_ext_ # optional: table-name prefix this extension
|
||||
# # owns under its own MetaData/migration chain.
|
||||
# # Keeps `alembic revision --autogenerate`, run
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user