diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 108eceed..805ec23b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -31,11 +31,14 @@ jobs: - name: Install just uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - - name: Setup dify-plugin-cli - run: ./scripts/setup-dify-plugin-cli.sh - - - name: Install dependencies - run: just dev - - - name: Run tests - run: just test + - name: Install dependencies and run tests + run: | + uv sync --locked + just coverage + + - name: Upload test reports + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: test-reports-${{ matrix.python-version }} + path: reports/ diff --git a/.gitignore b/.gitignore index 07f6334c..3aec53cc 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,10 @@ dist/ # Unit test .pytest_cache/ +/.hypothesis/ +/.coverage +/.coverage.* +/reports/ # Environments .venv/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9e895530..cba0ff2b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -35,7 +35,8 @@ source .venv/bin/activate The repository uses [`uv`](https://docs.astral.sh/uv/) for dependency and virtual environment management. The default development environment includes -`pytest`, `pytest-cov`, `pytest-mock`, `pytest-xprocess`, `ruff`, `ty`, and +`pytest`, `pytest-cov`, `pytest-mock`, `pytest-httpserver`, `pytest-timeout`, +`hypothesis`, `trustme`, `ruff`, `ty`, and [`prek`](https://prek.j178.dev/). ### Git Hooks @@ -65,7 +66,7 @@ Use `just` by default. For targeted work, direct tool usage is still fine: ```bash uv run ruff check src/dify_plugin/path.py -uv run pytest tests/path/test_file.py -k keyword +just test tests/path/test_file.py -k keyword uv run prek run -a ``` @@ -77,6 +78,7 @@ Use these commands for normal development: - `just lint`: run `just fmt`, then `uv run ruff check --fix` - `just check`: run `uv lock --check`, `ruff format --check --diff`, and `ruff check` - `just test`: run the SDK tests +- `just coverage`: run tests with subprocess and branch coverage, then enforce the coverage gates - `just build`: build source and wheel distributions - `just docs`: generate schema documentation into `.mkdocs/docs/schema.md` - `just clean`: remove local build, test, and lint artifacts @@ -84,8 +86,6 @@ Use these commands for normal development: Notes: - `just check` is the non-mutating validation entrypoint used by PR checks and Git hooks. -- Integration tests that need the Dify plugin CLI are skipped when the binary is unavailable. -- CI installs the CLI with [`scripts/setup-dify-plugin-cli.sh`](scripts/setup-dify-plugin-cli.sh) before running tests. - If you change dependencies, refresh and commit [`uv.lock`](uv.lock) before opening a pull request. Run `just check` for every change, plus `just test` for behavior changes, `just build` for packaging changes, and `just docs` for schema changes. @@ -96,9 +96,7 @@ Pull requests targeting `main` currently run these checks: 1. PR title validation with `amannn/action-semantic-pull-request` 2. `just check` on Python 3.12, including `uv.lock` freshness validation -3. `just test` on Python 3.12 and 3.13 through the reusable test workflow - -The test workflow installs the Dify plugin CLI, runs `just dev`, and then runs `just test`. +3. `just coverage` on Python 3.12 and 3.13 through the reusable test workflow Pushes to `main` also run the MkDocs workflow. It runs `just docs` on Python 3.12 and deploys `.mkdocs` to GitHub Pages. diff --git a/justfile b/justfile index 4febb386..e289a438 100644 --- a/justfile +++ b/justfile @@ -15,8 +15,16 @@ check: uv run ruff format --check --diff uv run ruff check -test: - uv run pytest +[positional-arguments] +test *args: + uv run python -m tests "$@" + +coverage: + uv run python -m tests --cov \ + --cov-report=json:reports/coverage.json \ + --cov-report=xml:reports/coverage.xml \ + --junitxml=reports/junit.xml + uv run python scripts/check_coverage.py reports/coverage.json build: uv build --no-create-gitignore --no-sources diff --git a/pyproject.toml b/pyproject.toml index dcc91ba8..dd9200eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,12 +35,15 @@ build-backend = 'uv_build' [dependency-groups] dev = [ + "hypothesis>=6.169.0", 'prek', 'pytest', 'pytest-cov', + "pytest-httpserver>=1.2.0", 'pytest-mock', - 'pytest-xprocess', + "pytest-timeout>=2.4.0", 'ruff', + "trustme>=1.2.1", 'ty', ] @@ -50,6 +53,27 @@ default-groups = ['dev'] [tool.pytest] minversion = '9.0' testpaths = ['tests'] +addopts = ['--strict-config', '--strict-markers'] +# pytest-timeout registers a string ini value, even with pytest's native TOML config. +timeout = '60' +xfail_strict = true +markers = ['integration: runs a real plugin process or local protocol service'] +filterwarnings = [ + 'error::ResourceWarning', + 'error::pytest.PytestUnhandledThreadExceptionWarning', + 'error::pytest.PytestUnraisableExceptionWarning', +] + +[tool.coverage.run] +branch = true +source_pkgs = ['dify_plugin'] +# SDK threads are monkey-patched greenlets; a second thread tracer drops events. +concurrency = ['gevent'] +patch = ['subprocess'] +sigterm = true + +[tool.coverage.paths] +source = ['src/dify_plugin', '*/site-packages/dify_plugin'] [tool.ruff] preview = true diff --git a/tests/__mock_server/__init__.py b/scripts/__init__.py similarity index 100% rename from tests/__mock_server/__init__.py rename to scripts/__init__.py diff --git a/scripts/check_coverage.py b/scripts/check_coverage.py new file mode 100644 index 00000000..39193fe3 --- /dev/null +++ b/scripts/check_coverage.py @@ -0,0 +1,58 @@ +import json +import sys +from pathlib import Path + +GATES = ( + ("SDK", ("src/dify_plugin/",), 85, 75), + ( + "runtime", + ("src/dify_plugin/core/runtime.py", "src/dify_plugin/core/server/"), + 90, + 85, + ), + ( + "HTTP", + ( + "src/dify_plugin/interfaces/model/openai_compatible/", + "src/dify_plugin/file/", + "src/dify_plugin/invocations/file.py", + ), + 90, + 85, + ), +) + + +def check_coverage(report: dict) -> list[str]: + if not report["meta"]["branch_coverage"]: + return ["Branch coverage was not collected"] + failures = [] + for name, prefixes, minimum_lines, minimum_branches in GATES: + summaries = [ + data["summary"] + for path, data in report["files"].items() + if path.startswith(prefixes) + ] + if not summaries: + failures.append(f"{name}: no matching source files in coverage report") + continue + for metric, covered_key, total_key, minimum in ( + ("lines", "covered_lines", "num_statements", minimum_lines), + ("branches", "covered_branches", "num_branches", minimum_branches), + ): + covered = sum(summary[covered_key] for summary in summaries) + total = sum(summary[total_key] for summary in summaries) + percent = 100 * covered / total if total else 0 + sys.stdout.write(f"{name} {metric}: {percent:.2f}% (minimum {minimum}%)\n") + if percent < minimum: + failures.append( + f"{name} {metric}: {covered}/{total} ({percent:.4f}%) < {minimum}%" + ) + return failures + + +if __name__ == "__main__": + failures = check_coverage(json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))) + if failures: + sys.stderr.write("\n".join(failures) + "\n") + raise SystemExit(bool(failures)) diff --git a/src/dify_plugin/core/documentation/generator.py b/src/dify_plugin/core/documentation/generator.py index 927d5942..da1409c8 100644 --- a/src/dify_plugin/core/documentation/generator.py +++ b/src/dify_plugin/core/documentation/generator.py @@ -38,7 +38,6 @@ def __init__(self) -> None: self._processed_types: set[type] = set() self._field_descriptions: dict[tuple[type, str], str] = {} self._schema_descriptions: dict[type, str] = {} - self._processed_field_types: set[type] = set() self._type_to_schema: dict[type, object] = {} self._type_blocks: dict[type, int] = {} self._blocks: list[list] = [] @@ -186,15 +185,15 @@ def _preprocess_schemas(self, schemas: list) -> None: # First pass: collect all field descriptions for schema in schemas: cls = schema.cls - if not issubclass(cls, BaseModel): - continue - # Store schema description if cls not in self._schema_descriptions or len(schema.description) > len( self._schema_descriptions[cls], ): self._schema_descriptions[cls] = schema.description + if not issubclass(cls, BaseModel): + continue + # Store field descriptions outside_reference_fields = ( getattr(schema, "outside_reference_fields", {}) or {} @@ -381,12 +380,6 @@ def _write_schema_doc(self, f: TextIO, type_: type) -> None: if field_type is None: continue - # Skip if we've already processed this field type - if isinstance(field_type, type) and issubclass(field_type, BaseModel): - if field_type in self._processed_field_types: - continue - self._processed_field_types.add(field_type) - # Skip if we've already processed this field field_key = (field_type, field_name) if field_key in processed_fields: diff --git a/src/dify_plugin/core/plugin_executor.py b/src/dify_plugin/core/plugin_executor.py index 0a0560c4..076fa19c 100644 --- a/src/dify_plugin/core/plugin_executor.py +++ b/src/dify_plugin/core/plugin_executor.py @@ -694,42 +694,45 @@ def invoke_endpoint( except Exception as e: response = Response(f"Internal Server Error: {e!s}", status=500) - # check if response is a generator - if isinstance(response.response, Generator): - # return headers - yield { - "status": response.status_code, - "headers": dict(response.headers.items()), - } - - for chunk in response.response: - if isinstance(chunk, bytes | bytearray | memoryview): - yield {"result": binascii.hexlify(chunk).decode()} - else: - yield {"result": binascii.hexlify(chunk.encode("utf-8")).decode()} - else: - result = { - "status": response.status_code, - "headers": dict(response.headers.items()), - } + with response: + # check if response is a generator + if isinstance(response.response, Generator): + # return headers + yield { + "status": response.status_code, + "headers": dict(response.headers.items()), + } - if isinstance(response.response, bytes | bytearray | memoryview): - result["result"] = binascii.hexlify(response.response).decode() - elif isinstance(response.response, str): - result["result"] = binascii.hexlify( - response.response.encode("utf-8"), - ).decode() - elif isinstance(response.response, Iterable): - result["result"] = "" for chunk in response.response: if isinstance(chunk, bytes | bytearray | memoryview): - result["result"] += binascii.hexlify(chunk).decode() + yield {"result": binascii.hexlify(chunk).decode()} else: - result["result"] += binascii.hexlify( - chunk.encode("utf-8"), - ).decode() - - yield result + yield { + "result": binascii.hexlify(chunk.encode("utf-8")).decode() + } + else: + result = { + "status": response.status_code, + "headers": dict(response.headers.items()), + } + + if isinstance(response.response, bytes | bytearray | memoryview): + result["result"] = binascii.hexlify(response.response).decode() + elif isinstance(response.response, str): + result["result"] = binascii.hexlify( + response.response.encode("utf-8"), + ).decode() + elif isinstance(response.response, Iterable): + result["result"] = "" + for chunk in response.response: + if isinstance(chunk, bytes | bytearray | memoryview): + result["result"] += binascii.hexlify(chunk).decode() + else: + result["result"] += binascii.hexlify( + chunk.encode("utf-8"), + ).decode() + + yield result def _get_oauth_provider_instance( self, @@ -966,14 +969,15 @@ def dispatch_trigger_event( subscription=subscription, request=original_request, ) - return TriggerDispatchResponse( - user_id=dispatch_result.user_id, - events=dispatch_result.events, - response=binascii.hexlify( - data=serialize_response(response=dispatch_result.response), - ).decode(), - payload=dispatch_result.payload, - ) + with dispatch_result.response: + return TriggerDispatchResponse( + user_id=dispatch_result.user_id, + events=dispatch_result.events, + response=binascii.hexlify( + data=serialize_response(response=dispatch_result.response), + ).decode(), + payload=dispatch_result.payload, + ) def subscribe_trigger( self, diff --git a/src/dify_plugin/core/plugin_registration.py b/src/dify_plugin/core/plugin_registration.py index 723e52fb..7ce99071 100644 --- a/src/dify_plugin/core/plugin_registration.py +++ b/src/dify_plugin/core/plugin_registration.py @@ -324,16 +324,32 @@ def _resolve_trigger_providers(self) -> None: parent_type=Trigger, ) + constructor_source = source subscription_constructor_cls_candidates = load_multi_subclasses_from_source( module_name=module_source, script_path=_source_to_script_path(source), parent_type=TriggerSubscriptionConstructor, ) + constructor = provider.subscription_constructor + if ( + not subscription_constructor_cls_candidates + and constructor + and constructor.extra + and constructor.extra.python.source != source + ): + constructor_source = constructor.extra.python.source + subscription_constructor_cls_candidates = ( + load_multi_subclasses_from_source( + module_name=_source_to_module_name(constructor_source), + script_path=_source_to_script_path(constructor_source), + parent_type=TriggerSubscriptionConstructor, + ) + ) if len(subscription_constructor_cls_candidates) > 1: msg = ( "Multiple TriggerSubscriptionConstructor subclasses found " - f"in {source}." + f"in {constructor_source}." " Only a single implementation is supported." ) raise ValueError(msg) @@ -350,7 +366,7 @@ def _resolve_trigger_providers(self) -> None: ): msg = ( "Trigger subscription constructor configuration declared " - f"but no implementation found in {source}." + f"but no implementation found in {constructor_source}." ) raise ValueError(msg) diff --git a/src/dify_plugin/core/runtime.py b/src/dify_plugin/core/runtime.py index 71df0413..5a333991 100644 --- a/src/dify_plugin/core/runtime.py +++ b/src/dify_plugin/core/runtime.py @@ -396,28 +396,29 @@ def _full_duplex_backwards_invoke( data_type: type[T], data: dict, ) -> Generator[T, None, None]: - if not self.session: + session = self.session + if not session: msg = "current tool runtime does not support backwards invoke" raise Exception(msg) - self.session.writer.session_message( - session_id=self.session.session_id, - data=self.session.writer.stream_invoke_object( - data={ - "type": type.value, - "backwards_request_id": backwards_request_id, - "request": data, - } - ), - ) - def filter(data: PluginInStream) -> bool: # ruff:ignore[builtin-variable-shadowing] return ( - data.event == PluginInStreamEvent.BackwardInvocationResponse + data.session_id == session.session_id + and data.event == PluginInStreamEvent.BackwardInvocationResponse and data.data.get("backwards_request_id") == backwards_request_id ) - with self.session.reader.read(filter) as reader: + with session.reader.read(filter) as reader: + session.writer.session_message( + session_id=session.session_id, + data=session.writer.stream_invoke_object( + data={ + "type": type.value, + "backwards_request_id": backwards_request_id, + "request": data, + } + ), + ) yield from self._line_converter_wrapper( reader.read(timeout_for_round=1), data_type ) diff --git a/src/dify_plugin/core/server/__base/filter_reader.py b/src/dify_plugin/core/server/__base/filter_reader.py index e29f1935..3b11cf85 100644 --- a/src/dify_plugin/core/server/__base/filter_reader.py +++ b/src/dify_plugin/core/server/__base/filter_reader.py @@ -20,6 +20,7 @@ def __init__( self.filter = filter self.queue = Queue() self.close_callback = close_callback + self._closed = False @overload def read( @@ -32,7 +33,7 @@ def read(self) -> Generator[PluginInStream, None, None]: ... def read( self, timeout_for_round: float | None = None ) -> Generator[PluginInStream | None, None, None]: - while True: + while not self._closed or not self.queue.empty(): try: data = self.queue.get(timeout=timeout_for_round) if data is None: @@ -45,13 +46,19 @@ def read( break def close(self) -> None: - if self.close_callback: - self.close_callback() + if self._closed: + return - self.queue.put(None) + self._closed = True + try: + if self.close_callback: + self.close_callback() + finally: + self.queue.put(None) def write(self, data: PluginInStream) -> None: - self.queue.put(data) + if not self._closed: + self.queue.put(data) def __enter__(self) -> Self: return self diff --git a/src/dify_plugin/core/server/__base/request_reader.py b/src/dify_plugin/core/server/__base/request_reader.py index 86186b16..b6012477 100644 --- a/src/dify_plugin/core/server/__base/request_reader.py +++ b/src/dify_plugin/core/server/__base/request_reader.py @@ -20,6 +20,7 @@ def __init__(self) -> None: # Convert class variables to instance variables to avoid global lock contention self.lock = threading.Lock() self.readers: list[FilterReader] = [] + self._closed = False @abstractmethod def _read_stream(self) -> Generator["PluginInStream", None, None]: @@ -29,14 +30,18 @@ def _read_stream(self) -> Generator["PluginInStream", None, None]: raise NotImplementedError def event_loop(self) -> None: - # read line by line - while True: - try: - for line in self._read_stream(): - self._process_line(line) - except Exception: - logger.exception("Error in event loop") - time.sleep(0.01) # Prevent high CPU usage after failures + try: + while not self._closed: + try: + for line in self._read_stream(): + if self._closed: + return + self._process_line(line) + except Exception: + logger.exception("Error in event loop") + time.sleep(0.01) # Prevent high CPU usage after failures + finally: + self.close() def _process_line(self, data: "PluginInStream") -> None: session_id = data.session_id @@ -94,6 +99,7 @@ def close(self) -> None: """ self.lock.acquire() try: + self._closed = True readers_to_close = self.readers.copy() self.readers.clear() finally: diff --git a/src/dify_plugin/core/server/__base/response_writer.py b/src/dify_plugin/core/server/__base/response_writer.py index 6456bb9e..23663fe9 100644 --- a/src/dify_plugin/core/server/__base/response_writer.py +++ b/src/dify_plugin/core/server/__base/response_writer.py @@ -39,8 +39,8 @@ def put( StreamOutputMessage( event=event, session_id=session_id, data=data ).model_dump_json() + + "\n\n" ) - self.write("\n\n") def error( self, session_id: str | None = None, data: dict | BaseModel | None = None diff --git a/src/dify_plugin/core/server/io_server.py b/src/dify_plugin/core/server/io_server.py index 10add4f7..dd6a9819 100644 --- a/src/dify_plugin/core/server/io_server.py +++ b/src/dify_plugin/core/server/io_server.py @@ -5,6 +5,7 @@ import traceback from abc import ABC, abstractmethod from concurrent.futures import ThreadPoolExecutor +from functools import partial from threading import Thread from dify_plugin.config.config import DifyPluginEnv @@ -12,6 +13,7 @@ from dify_plugin.core.server.__base.request_reader import RequestReader from dify_plugin.core.server.__base.response_writer import ResponseWriter from dify_plugin.core.server.serverless.request_reader import ServerlessRequestReader +from dify_plugin.core.server.serverless.response_writer import ServerlessResponseWriter from dify_plugin.core.server.stdio.request_reader import StdioRequestReader from dify_plugin.core.server.tcp.request_reader import TCPReaderWriter from dify_plugin.errors.model import InvokeError @@ -91,43 +93,52 @@ def _execute_request_in_thread( """ wrapper for _execute_request """ - # wait for the task to finish try: - self._execute_request( - session_id, - data, - reader, - writer, - conversation_id, - message_id, - app_id, - endpoint_id, - context, - ) - except Exception as e: - args = {"traceback": traceback.format_exc()} - if isinstance(e, InvokeError): - args["description"] = e.description - - if isinstance(reader, (TCPReaderWriter, ServerlessRequestReader)): - logger.exception( - "Unexpected error occurred when executing request", - exc_info=e, + # wait for the task to finish + try: + execute = partial( + self._execute_request, + session_id, + data, + reader, + writer, + conversation_id, + message_id, + app_id, + endpoint_id, + context, + ) + if isinstance(writer, ServerlessResponseWriter): + writer.execute(execute) + else: + execute() + except Exception as e: + args = {"traceback": traceback.format_exc()} + if isinstance(e, InvokeError): + args["description"] = e.description + + if isinstance(reader, (TCPReaderWriter, ServerlessRequestReader)): + logger.exception( + "Unexpected error occurred when executing request", + exc_info=e, + ) + + writer.session_message( + session_id=session_id, + data=writer.stream_error_object( + data={ + "error_type": type(e).__name__, + "message": str(e), + "args": args, + } + ), ) writer.session_message( - session_id=session_id, - data=writer.stream_error_object( - data={ - "error_type": type(e).__name__, - "message": str(e), - "args": args, - } - ), + session_id=session_id, data=writer.stream_end_object() ) - - writer.session_message(session_id=session_id, data=writer.stream_end_object()) - writer.done() + finally: + writer.done() def _heartbeat(self) -> None: """ diff --git a/src/dify_plugin/core/server/serverless/request_reader.py b/src/dify_plugin/core/server/serverless/request_reader.py index 4ac64a14..bd947cbe 100644 --- a/src/dify_plugin/core/server/serverless/request_reader.py +++ b/src/dify_plugin/core/server/serverless/request_reader.py @@ -8,6 +8,7 @@ import gevent.socket from flask import Flask, request from gevent.pywsgi import WSGIServer +from werkzeug.serving import BaseWSGIServer, make_server from dify_plugin.core.entities.plugin.io import ( PluginInStream, @@ -40,22 +41,33 @@ def __init__( self.threads = threads self.worker_connections = worker_connections self.max_single_connection_lifetime = max_single_connection_lifetime - self.request_queue = Queue[PluginInStream]() + self.request_queue = Queue[PluginInStream | None]() + self._shutdown = threading.Event() + self._http_server: BaseWSGIServer | None = None def _read_stream(self) -> Generator[PluginInStream, None, None]: """ Read request from http server """ - while True: - yield self.request_queue.get() + while (message := self.request_queue.get()) is not None: + yield message + + def close(self) -> None: + super().close() + self.request_queue.put(None) + self._shutdown.set() + with self.lock: + server, self._http_server = self._http_server, None + if server is not None: + server.shutdown() def handler(self) -> tuple[Generator[str, None, None], int] | tuple[str, int]: try: queue: Queue[str | None] = Queue() data = request.get_json() - event = PluginInStreamEvent.value_of(data["event"]) + writer = ServerlessResponseWriter(queue) plugin_in = PluginInStream( - event=event, + event=PluginInStreamEvent.value_of(data["event"]), session_id=data["session_id"], conversation_id=data.get("conversation_id"), message_id=data.get("message_id"), @@ -64,29 +76,39 @@ def handler(self) -> tuple[Generator[str, None, None], int] | tuple[str, int]: data=data["data"], context=data.get("context"), reader=self, - writer=ServerlessResponseWriter(queue), + writer=writer, ) self.request_queue.put(plugin_in) except Exception as e: return str(e), 500 def generate() -> Generator[str, None, None]: - refresh_time = time.time() - while True: - try: - response = queue.get(timeout=1) - except Empty: - if time.time() - refresh_time > self.max_single_connection_lifetime: - return - continue + try: + yield "" + refresh_time = time.time() + while True: + try: + response = queue.get(timeout=1) + except Empty: + if ( + time.time() - refresh_time + > self.max_single_connection_lifetime + ): + return + continue - if response is None: - return + if response is None: + return - refresh_time = time.time() - yield response + refresh_time = time.time() + yield response + finally: + writer.close() - return generate(), 200 + response = generate() + # Arm cleanup before Flask receives the response iterator. + next(response) + return response, 200 def health(self) -> tuple[str, int]: return "OK", 200 @@ -106,9 +128,22 @@ def _run(self) -> None: ) sys.stdout.write("* Server Worker: gevent.wsgi.WSGIServer\n") sys.stdout.flush() - server.serve_forever() + try: + with self.lock: + if self._closed: + return + server.start() + self._shutdown.wait() + finally: + server.stop() else: - self.app.run(host=self.host, port=self.port, threaded=True) + server = make_server(self.host, self.port, self.app, threaded=True) + with self.lock: + if self._closed: + server.server_close() + return + self._http_server = server + server.serve_forever() def launch(self) -> None: """ diff --git a/src/dify_plugin/core/server/serverless/response_writer.py b/src/dify_plugin/core/server/serverless/response_writer.py index 78c88b7b..fd1b9191 100644 --- a/src/dify_plugin/core/server/serverless/response_writer.py +++ b/src/dify_plugin/core/server/serverless/response_writer.py @@ -1,4 +1,9 @@ -from queue import Queue +from collections.abc import Callable +from contextvars import copy_context +from queue import Empty, Queue +from threading import Lock + +from gevent import Greenlet from dify_plugin.core.server.__base.response_writer import ResponseWriter @@ -10,9 +15,41 @@ class ServerlessResponseWriter(ResponseWriter): def __init__(self, queue: Queue) -> None: self.q = queue + self.closed = False + self._lock = Lock() + self._task: Greenlet | None = None + + def execute(self, callback: Callable[[], None]) -> None: + with self._lock: + if self.closed: + return + task = self._task = Greenlet.spawn(copy_context().run, callback) + try: + task.get() + finally: + with self._lock: + self._task = None + task.kill(block=False) def write(self, data: bytes) -> None: - self.q.put(data) + with self._lock: + if not self.closed: + self.q.put(data) def done(self) -> None: - self.q.put(None) + with self._lock: + if not self.closed: + self.closed = True + self.q.put(None) + + def close(self) -> None: + with self._lock: + self.closed = True + task = self._task + while True: + try: + self.q.get_nowait() + except Empty: + break + if task is not None: + task.kill(block=False) diff --git a/src/dify_plugin/core/server/stdio/request_reader.py b/src/dify_plugin/core/server/stdio/request_reader.py index 2c7b728c..e15772c9 100644 --- a/src/dify_plugin/core/server/stdio/request_reader.py +++ b/src/dify_plugin/core/server/stdio/request_reader.py @@ -27,10 +27,11 @@ def _read_async(self) -> bytes: def _read_stream(self) -> Generator[PluginInStream, None, None]: buffer = BytesIO() - while True: + while not self._closed: data = self._read_async() if not data: - continue + self.close() + return # Write new data to buffer buffer.write(data) diff --git a/src/dify_plugin/core/server/tcp/request_reader.py b/src/dify_plugin/core/server/tcp/request_reader.py index 21fae0f0..429ca7dd 100644 --- a/src/dify_plugin/core/server/tcp/request_reader.py +++ b/src/dify_plugin/core/server/tcp/request_reader.py @@ -5,7 +5,7 @@ import socket as native_socket import time from collections.abc import Callable, Generator -from threading import Lock +from threading import RLock from typing import Any from gevent import sleep @@ -46,20 +46,28 @@ def __init__( self.reconnect_timeout = reconnect_timeout self.alive = False self.on_connected = on_connected - self.opt_lock = Lock() + self.opt_lock = RLock() + self._connecting = False + self._reconnect_pending = False # handle SIGINT to exit the program smoothly due to the gevent limitation signal.signal(signal.SIGINT, lambda *_args, **_kwargs: os._exit(0)) def launch(self) -> None: """Launch the connection""" - self._launch() + with self.opt_lock: + if self.alive: + self.sock.close() + self.alive = False + self._closed = False + self._launch() def close(self) -> None: """Close the connection""" + super().close() if self.alive: - self.sock.close() self.alive = False + self.sock.close() def _write_to_sock(self, data: bytes) -> int: """Write data to the socket""" @@ -76,6 +84,9 @@ def _write_data(self, data: str) -> None: while data_bytes: try: sent = self._write_to_sock(data_bytes) + if sent == 0: + msg = "Connection is closed" + raise ConnectionError(msg) data_bytes = data_bytes[sent:] except BlockingIOError as error: if error.errno != errno.EAGAIN: @@ -85,32 +96,51 @@ def _write_data(self, data: str) -> None: self.sock.sendall(data.encode()) def write(self, data: str) -> None: - if not self.alive: - msg = "connection is dead" - raise Exception(msg) + with self.opt_lock: + if not self.alive: + msg = "connection is dead" + raise Exception(msg) - try: - self._write_data(data) - except Exception: - logger.exception("Failed to write data") - self._launch() + sock = self.sock + try: + self._write_data(data) + except BaseException as error: + logger.exception("Failed to write data") + sock.close() + self.alive = False + # Registration failures belong to the outer connection retry loop. + if not self._closed and not self._connecting: + if isinstance(error, Exception): + self._launch() + return + # Reconnect in the reader so cancellation can finish now. + self._reconnect_pending = True + raise def done(self) -> None: pass def _launch(self) -> None: """Connect to the target, try to reconnect if failed""" - attempts = 0 - while attempts < self.reconnect_attempts: + with self.opt_lock: + if self.alive or self._closed: + return + self._reconnect_pending = False + self._connecting = True try: - self._connect() - break - except Exception: - attempts += 1 - if attempts >= self.reconnect_attempts: - raise + attempts = 0 + while attempts < self.reconnect_attempts and not self._closed: + try: + self._connect() + break + except Exception: + attempts += 1 + if attempts >= self.reconnect_attempts: + raise - time.sleep(self.reconnect_timeout) + time.sleep(self.reconnect_timeout) + finally: + self._connecting = False def _connect(self) -> None: """Connect to the target""" @@ -126,23 +156,33 @@ def _connect(self) -> None: def _connect_once(self) -> None: if native_socket.socket is gevent_socket.socket: - self.sock = gevent_socket.create_connection((self.host, self.port)) + sock = gevent_socket.create_connection((self.host, self.port)) else: - self.sock = native_socket.create_connection((self.host, self.port)) + sock = native_socket.create_connection((self.host, self.port)) + if self._closed: + sock.close() + return + self.sock = sock self.alive = True handshake_message = InitializeMessage( type=InitializeMessage.Type.HANDSHAKE, data=InitializeMessage.Key(key=self.key).model_dump(), ) - self.sock.sendall(handshake_message.model_dump_json().encode() + b"\n") - logger.info("\033[32mConnected to %s:%s\033[0m", self.host, self.port) - if self.on_connected: - self.on_connected() + try: + sock.sendall(handshake_message.model_dump_json().encode() + b"\n") + logger.info("\033[32mConnected to %s:%s\033[0m", self.host, self.port) + if self.on_connected: + self.on_connected() + except BaseException: + self.alive = False + sock.close() + raise logger.info("Sent key to %s:%s", self.host, self.port) def _read_data(self) -> bytes | None: - ready_to_read, _, _ = select([self.sock], [], [], 1) - if not ready_to_read: + sock = self.sock + ready_to_read, _, _ = select([sock], [], [], 1) + if not ready_to_read or sock is not self.sock: return None try: data = self._recv_from_sock(1048576) @@ -161,20 +201,42 @@ def _read_data(self) -> bytes | None: def _read_stream(self) -> Generator[PluginInStream, None, None]: """Read data from the target""" buffer = b"" - while self.alive: + source_socket = None + while not self._closed: + if not self.alive: + with self.opt_lock: + if not self.alive: + if not self._reconnect_pending: + self.close() + return + self._launch() + + sock = self.sock + if sock is not source_socket: + buffer = b"" + source_socket = sock try: data = self._read_data() except Exception: - logger.exception( - "\033[31mFailed to read data from %s:%s\033[0m", - self.host, - self.port, - ) - self.alive = False - time.sleep(self.reconnect_timeout) - self._launch() + with self.opt_lock: + if self._closed: + return + if sock is self.sock: + logger.exception( + "\033[31mFailed to read data from %s:%s\033[0m", + self.host, + self.port, + ) + sock.close() + self.alive = False + time.sleep(self.reconnect_timeout) + self._launch() + buffer = b"" continue + if sock is not self.sock: + buffer = b"" + continue if not data: continue @@ -182,13 +244,12 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]: # process line by line and keep the last line if it is not complete lines = buffer.split(b"\n") - if len(lines) == 0: - continue - buffer = lines[-1] lines = lines[:-1] for line in lines: + if not line.strip(): + continue try: data = TypeAdapter(dict[str, Any]).validate_json(line) chunk = PluginInStream( diff --git a/src/dify_plugin/core/utils/position_helper.py b/src/dify_plugin/core/utils/position_helper.py index 5bcb80ac..7042459d 100644 --- a/src/dify_plugin/core/utils/position_helper.py +++ b/src/dify_plugin/core/utils/position_helper.py @@ -19,7 +19,7 @@ def get_position_map( The return value. """ position_file_name = str(pathlib.Path(folder_path) / file_name) - positions = load_yaml_file(position_file_name, ignore_error=True) + positions = load_yaml_file(position_file_name, ignore_error=True) or {} position_map = {} index = 0 for name in positions: diff --git a/src/dify_plugin/entities/model/provider.py b/src/dify_plugin/entities/model/provider.py index d35695ab..ee939c76 100644 --- a/src/dify_plugin/entities/model/provider.py +++ b/src/dify_plugin/entities/model/provider.py @@ -237,6 +237,7 @@ def to_simple_provider(self) -> SimpleProviderEntity: @model_validator(mode="before") @classmethod def validate_models(cls, values: dict[str, object]) -> dict[str, object]: + values = dict(values) value = values.get("models", {}) if not isinstance(value, dict): msg = "models should be a glob path list" diff --git a/src/dify_plugin/entities/tool.py b/src/dify_plugin/entities/tool.py index 75c5f79b..f981b0f9 100644 --- a/src/dify_plugin/entities/tool.py +++ b/src/dify_plugin/entities/tool.py @@ -294,6 +294,7 @@ class ToolProviderConfiguration(BaseModel): @model_validator(mode="before") @classmethod def validate_credentials_schema(cls, data: dict) -> dict: + data = dict(data) original_credentials_for_provider: dict[str, dict] = data.get( "credentials_for_provider", {}, diff --git a/src/dify_plugin/integration/run.py b/src/dify_plugin/integration/run.py index bb0db3fb..5ad46119 100644 --- a/src/dify_plugin/integration/run.py +++ b/src/dify_plugin/integration/run.py @@ -69,68 +69,69 @@ def __init__( self.config = config self.plugin_package_path = plugin_package_path self.extra_args = extra_args or [] - self.resources_need_to_be_cleaned = [] - - # create pipe to communicate with the plugin - self.stdout_pipe_read, self.stdout_pipe_write = os.pipe() - self.stderr_pipe_read, self.stderr_pipe_write = os.pipe() - self.stdin_pipe_read, self.stdin_pipe_write = os.pipe() - - # stdin write lock + self.resources_need_to_be_cleaned: list[str] = [] + self.stdout_pipe_read = self.stdout_pipe_write = -1 + self.stderr_pipe_read = self.stderr_pipe_write = -1 + self.stdin_pipe_read = self.stdin_pipe_write = -1 self.stdin_write_lock = Lock() - - # setup stop flag self.stop_flag = False self.stop_flag_lock = Lock() - - logger.info("Running plugin from %s", plugin_package_path) - - # check if plugin is a directory - if pathlib.Path(plugin_package_path).is_dir(): - logger.info("plugin source directory detected, building plugin") - with tempfile.TemporaryDirectory(delete=False) as temp_dir: - output_path = str(pathlib.Path(temp_dir) / "plugin.difypkg") + self.q: dict[str, Queue[PluginGenericResponse | Exception | None]] = {} + self.q_lock = Lock() + self.ready_semaphore = Semaphore(0) + self._close_complete = threading.Event() + self.process: subprocess.Popen[bytes] | None = None + self.stdout_reader: threading.Thread | None = None + self._stderr_reader: threading.Thread | None = None + self._temporary_directory: tempfile.TemporaryDirectory[str] | None = None + self._reader_error: Exception | None = None + + # Every startup failure must release the process and temporary package. + try: # ruff:ignore[too-many-statements-in-try-clause] + self.stdout_pipe_read, self.stdout_pipe_write = os.pipe() + self.stderr_pipe_read, self.stderr_pipe_write = os.pipe() + self.stdin_pipe_read, self.stdin_pipe_write = os.pipe() + logger.info("Running plugin from %s", plugin_package_path) + if pathlib.Path(plugin_package_path).is_dir(): + logger.info("plugin source directory detected, building plugin") + self._temporary_directory = tempfile.TemporaryDirectory() + self.resources_need_to_be_cleaned.append(self._temporary_directory.name) + output_path = str( + pathlib.Path(self._temporary_directory.name) / "plugin.difypkg" + ) self._build_plugin(plugin_package_path, output_path) self.plugin_package_path = output_path logger.info("Plugin built in %s", self.plugin_package_path) - self.resources_need_to_be_cleaned.append(temp_dir) - - self.process = subprocess.Popen( # ruff:ignore[subprocess-without-shell-equals-true] - [ - self.config.dify_cli_path, - "plugin", - "run", - self.plugin_package_path, - "--response-format", - "json", - *self.extra_args, - ], - stdout=self.stdout_pipe_write, - stderr=self.stderr_pipe_write, - stdin=self.stdin_pipe_read, - ) - - logger.info("Plugin process created with pid %s", self.process.pid) - - # wait for plugin to be ready - self.ready_semaphore = Semaphore(0) - - # create a thread to read the stdout and stderr - self.stdout_reader = threading.Thread( - target=self._message_reader, - args=(self.stdout_pipe_read,), - ) - self.stdout_reader.start() - self.q = dict[str, Queue[PluginGenericResponse | None]]() - self.q_lock = Lock() - - # wait for the plugin to be ready with timeout - if not self.ready_semaphore.acquire(timeout=30): # 30 seconds timeout - msg = "Plugin failed to start within 30 seconds" - raise TimeoutError(msg) - - logger.info("Plugin ready") + self.process = subprocess.Popen( # ruff:ignore[subprocess-without-shell-equals-true] + [ + self.config.dify_cli_path, + "plugin", + "run", + self.plugin_package_path, + "--response-format", + "json", + *self.extra_args, + ], + stdin=self.stdin_pipe_read, + stdout=self.stdout_pipe_write, + stderr=self.stderr_pipe_write, + ) + logger.info("Plugin process created with pid %s", self.process.pid) + self.stdout_reader = threading.Thread( + target=self._message_reader, args=(self.stdout_pipe_read,) + ) + self._stderr_reader = threading.Thread(target=self._drain_stderr) + self.stdout_reader.start() + self._stderr_reader.start() + + if not self.ready_semaphore.acquire(timeout=30): + msg = "Plugin failed to start within 30 seconds" + raise TimeoutError(msg) + logger.info("Plugin ready") + except BaseException: + self._close() + raise def _build_plugin(self, package_path: str, output_path: str) -> None: # build plugin @@ -147,70 +148,88 @@ def _build_plugin(self, package_path: str, output_path: str) -> None: logger.info(output.decode("utf-8")) def _close(self) -> None: - with self.stop_flag_lock: - if self.stop_flag: - return - - # stop the plugin + with self.stop_flag_lock, self.q_lock: + first_close = not self.stop_flag self.stop_flag = True + if first_close: + for queue in self.q.values(): + queue.put(self._reader_error or PluginStoppedError()) - # send signal SIGTERM to the plugin, so it can exit gracefully - # do collect garbage like removing temporary files - os.kill(self.process.pid, signal.SIGTERM) - - # wait for the plugin to exit - self.process.wait() + if not first_close: + # The owner may be joining this reader while it finishes. + if threading.current_thread() in {self.stdout_reader, self._stderr_reader}: + return + self._close_complete.wait() + return - # close the pipes - os.close(self.stdout_pipe_write) - os.close(self.stderr_pipe_write) - os.close(self.stdin_pipe_read) + try: + if self.process is not None: + if self.process.poll() is None: + self.process.send_signal(signal.SIGTERM) + self.process.wait() + + for thread in (self.stdout_reader, self._stderr_reader): + if ( + thread is not None + and thread.ident is not None + and thread is not threading.current_thread() + ): + thread.join() + + for pipe in ( + self.stdout_pipe_read, + self.stdout_pipe_write, + self.stderr_pipe_read, + self.stderr_pipe_write, + self.stdin_pipe_read, + self.stdin_pipe_write, + ): + if pipe != -1: + os.close(pipe) + for resource in self.resources_need_to_be_cleaned: + shutil.rmtree(resource) + if self._temporary_directory is not None: + self._temporary_directory.cleanup() + finally: + self._close_complete.set() def _read_async(self, fd: int) -> bytes: ready, _, _ = select.select([fd], [], [], 0.1) if not ready: + if self.process is not None and self.process.poll() is not None: + raise PluginStoppedError return b"" - - # read data from stdin using os.read in 64KB chunks. - # the OS buffer for stdin is usually 64KB, so using a larger value - # doesn't make sense. - b = os.read(fd, 65536) - if not b: + data = os.read(fd, 65536) + if not data: raise PluginStoppedError - return b + return data - def _message_reader(self, pipe: int) -> None: - # create a scanner to read the message line by line - """Read messages line by line from the pipe.""" - buffer = b"" + def _drain_stderr(self) -> None: try: while True: - try: - data = self._read_async(pipe) - except PluginStoppedError: - break + # Keep CLI logs outside subclass stdout parsing hooks. + PluginRunner._read_async(self, self.stderr_pipe_read) + except (PluginStoppedError, OSError): + return + def _message_reader(self, pipe: int) -> None: + buffer = b"" + try: # ruff:ignore[too-many-statements-in-try-clause] + while True: + data = self._read_async(pipe) if not data: time.sleep(0.01) continue - - buffer += data - - # if no b"\n" is in data, skip to the next iteration - if data.find(b"\n") == -1: - continue - - # process line by line and keep the last line if it is not complete - lines = buffer.split(b"\n") - buffer = lines[-1] - - lines = lines[:-1] + lines = (buffer + data).split(b"\n") + buffer = lines.pop() for raw_line in lines: line = raw_line.strip() - if not line: - continue - - self._publish_message(line.decode("utf-8")) + if line: + self._publish_message(line.decode("utf-8")) + except PluginStoppedError: + pass + except Exception as error: + self._reader_error = error finally: self._close() @@ -242,13 +261,11 @@ def _publish_message(self, message: str) -> None: self.q[parsed_message.invoke_id].put(parsed_message) def _write_to_pipe(self, data: bytes) -> None: - # split the data into chunks of 4096 bytes - chunks = [data[i : i + 4096] for i in range(0, len(data), 4096)] - # A lock is needed to avoid race conditions when multiple threads - # write to the pipe. with self.stdin_write_lock: - for chunk in chunks: - os.write(self.stdin_pipe_write, chunk) + remaining = memoryview(data) + while remaining: + written = os.write(self.stdin_pipe_write, remaining[:4096]) + remaining = remaining[written:] def invoke( self, @@ -269,8 +286,10 @@ def invoke( request=payload, ) - q = Queue[PluginGenericResponse | None]() + q = Queue[PluginGenericResponse | Exception | None]() with self.q_lock: + if self.stop_flag: + raise PluginStoppedError self.q[invoke_id] = q try: @@ -278,7 +297,9 @@ def invoke( self._write_to_pipe(request.model_dump_json().encode("utf-8") + b"\n") # wait for events - while message := q.get(): + while (message := q.get()) is not None: + if isinstance(message, Exception): + raise message if message.invoke_id == invoke_id: if message.type == ResponseType.PLUGIN_RESPONSE: yield response_type.model_validate(message.response) @@ -292,7 +313,7 @@ def invoke( raise ValueError(msg) finally: with self.q_lock: - del self.q[invoke_id] + self.q.pop(invoke_id, None) def __enter__(self) -> Self: return self @@ -304,6 +325,3 @@ def __exit__( traceback: TracebackType | None, ) -> None: self._close() - - for resource in self.resources_need_to_be_cleaned: - shutil.rmtree(resource) diff --git a/src/dify_plugin/interfaces/model/ai_model.py b/src/dify_plugin/interfaces/model/ai_model.py index 83cd9b00..2baaf358 100644 --- a/src/dify_plugin/interfaces/model/ai_model.py +++ b/src/dify_plugin/interfaces/model/ai_model.py @@ -120,7 +120,7 @@ def _transform_invoke_error(self, error: Exception) -> InvokeError: Returns: The return value. """ - provider_name = self.__class__.__module__.split(".")[-3] + provider_name = self.__class__.__module__.split(".")[-3:][0] for invoke_error, model_errors in self._invoke_error_mapping.items(): if isinstance(error, tuple(model_errors)): diff --git a/src/dify_plugin/interfaces/model/large_language_model.py b/src/dify_plugin/interfaces/model/large_language_model.py index 16f6bd49..c6c2e955 100644 --- a/src/dify_plugin/interfaces/model/large_language_model.py +++ b/src/dify_plugin/interfaces/model/large_language_model.py @@ -4,6 +4,7 @@ import time from abc import abstractmethod from collections.abc import Generator, Mapping +from contextlib import closing from typing import Literal from pydantic import ConfigDict, JsonValue @@ -550,21 +551,32 @@ def new_generator() -> Generator[LLMResultChunk, None, None]: yield first_chunk yield from response + processor = self._code_block_mode_stream_processor if ( first_chunk.delta.message.content and isinstance(first_chunk.delta.message.content, str) and first_chunk.delta.message.content.startswith("`") ): - return self._code_block_mode_stream_processor_with_backtick( + processor = self._code_block_mode_stream_processor_with_backtick + try: + processed_response = processor( model=model, prompt_messages=prompt_messages, input_generator=new_generator(), ) - return self._code_block_mode_stream_processor( - model=model, - prompt_messages=prompt_messages, - input_generator=new_generator(), - ) + except BaseException: + response.close() + raise + + def process_response() -> Generator[LLMResultChunk, None, None]: + with closing(response): + yield first_chunk + yield from processed_response + + result = process_response() + # Activate cleanup without consuming another provider chunk. + next(result) + return result return response diff --git a/src/dify_plugin/interfaces/model/openai_compatible/llm.py b/src/dify_plugin/interfaces/model/openai_compatible/llm.py index 22e4ead4..bcfb42f4 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/llm.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/llm.py @@ -8,6 +8,7 @@ from contextlib import suppress from decimal import Decimal from http import HTTPStatus +from io import IncrementalNewlineDecoder from typing import Any, cast import requests @@ -53,6 +54,22 @@ _plugin_config = DifyPluginEnv() +def _iter_sse_events(response: requests.Response) -> Generator[str]: + encoding = response.encoding or "utf-8" + if codecs.lookup(encoding).name == "utf-8": + encoding = "utf-8-sig" + decoder = IncrementalNewlineDecoder( + codecs.getincrementaldecoder(encoding)("replace"), translate=True + ) + pending = "" + for data in response.iter_content(chunk_size=512): + pending += decoder.decode(data) + *events, pending = pending.split("\n\n") + yield from events + pending += decoder.decode(b"", final=True) + yield from pending.split("\n\n") + + def _gen_tool_call_id() -> str: return f"chatcmpl-tool-{uuid.uuid4().hex!s}" @@ -465,7 +482,7 @@ def get_customizable_model_schema( elif credentials["mode"] == "completion": entity.model_properties[ModelPropertyKey.MODE] = LLMMode.COMPLETION.value else: - msg = f"Unknown completion type {credentials['completion_type']}" + msg = f"Unknown completion type {credentials['mode']}" raise ValueError( msg, ) @@ -715,13 +732,16 @@ def _handle_generate_stream_response( # delimiter for stream response, need unicode_escape delimiter = credentials.get("stream_mode_delimiter", "\n\n") delimiter = codecs.decode(delimiter, "unicode_escape") - for raw_chunk in response.iter_lines( - decode_unicode=True, delimiter=delimiter - ): + events = ( + _iter_sse_events(response) + if delimiter == "\n\n" + else response.iter_lines(decode_unicode=True, delimiter=delimiter) + ) + for raw_chunk in events: chunk = raw_chunk.strip() if not chunk or chunk.startswith(":"): continue - decoded_chunk = chunk.strip().removeprefix("data:").lstrip() + decoded_chunk = chunk.removeprefix("data:").lstrip() if decoded_chunk == "[DONE]": # Some provider returns "data: [DONE]" break @@ -731,7 +751,28 @@ def _handle_generate_stream_response( ) except ValidationError as error: msg = "Invalid JSON in the model response stream" - raise InvokeError(msg) from error + lines = chunk.split("\n") + if any( + not line.startswith(("data:", "event:", "id:", "retry:", ":")) + for line in lines + ): + raise InvokeError(msg) from error + data_lines = [ + line.removeprefix("data:").lstrip() + for line in lines + if line.startswith("data:") + ] + if not data_lines: + raise InvokeError(msg) from error + decoded_data = "\n".join(data_lines) + if decoded_data == "[DONE]": + break + try: + chunk_json = TypeAdapter(dict[str, Any]).validate_json( + decoded_data, + ) + except ValidationError as data_error: + raise InvokeError(msg) from data_error # handle the error here. for issue #11629 if chunk_json.get("error"): raise InvokeError(str(chunk_json["error"])) @@ -1199,11 +1240,14 @@ def _num_tokens_for_tool_property(self, key: str, value: dict) -> int: num_tokens += self._get_num_tokens_by_gpt2(str(field_value)) return num_tokens - def _num_tokens_for_enum_field(self, field_value: list[str]) -> int: + def _num_tokens_for_enum_field(self, field_value: list[Any]) -> int: num_tokens = 0 for enum_field in field_value: num_tokens += 3 - num_tokens += self._get_num_tokens_by_gpt2(enum_field) + value = ( + enum_field if isinstance(enum_field, str) else json.dumps(enum_field) + ) + num_tokens += self._get_num_tokens_by_gpt2(value) return num_tokens def _extract_response_tool_calls( diff --git a/src/dify_plugin/interfaces/model/openai_compatible/tts.py b/src/dify_plugin/interfaces/model/openai_compatible/tts.py index f36b350b..ee88657e 100644 --- a/src/dify_plugin/interfaces/model/openai_compatible/tts.py +++ b/src/dify_plugin/interfaces/model/openai_compatible/tts.py @@ -1,4 +1,5 @@ from collections.abc import Generator +from contextlib import closing from http import HTTPStatus import requests @@ -79,21 +80,19 @@ def _invoke( } # Make POST request - response = requests.post( + with requests.post( endpoint_url, headers=headers, json=payload, stream=True, timeout=(10, 300), - ) + ) as response: + if response.status_code != HTTPStatus.OK: + raise InvokeBadRequestError(response.text) - if response.status_code != HTTPStatus.OK: - raise InvokeBadRequestError(response.text) - - # Stream the audio data - for chunk in response.iter_content(chunk_size=4096): - if chunk: - yield chunk + for chunk in response.iter_content(chunk_size=4096): + if chunk: + yield chunk def validate_credentials(self, model: str, credentials: dict) -> None: """ @@ -111,7 +110,7 @@ def validate_credentials(self, model: str, credentials: dict) -> None: voice = self._get_model_default_voice(model, credentials) # Test with a simple text - next( + with closing( self._invoke( model=model, tenant_id="validate", @@ -119,7 +118,8 @@ def validate_credentials(self, model: str, credentials: dict) -> None: content_text="Test.", voice=voice, ) - ) + ) as stream: + next(stream) except Exception as ex: raise CredentialsValidateFailedError(str(ex)) from ex diff --git a/src/dify_plugin/invocations/model/llm.py b/src/dify_plugin/invocations/model/llm.py index 0e703451..33ca0c6e 100644 --- a/src/dify_plugin/invocations/model/llm.py +++ b/src/dify_plugin/invocations/model/llm.py @@ -17,6 +17,25 @@ ) +def merge_llm_result_chunk(result: LLMResult, chunk: LLMResultChunk) -> None: + content = chunk.delta.message.content + # Keep blocking results compatible with callers that expect a string. + if isinstance(content, str): + result.message.content = cast("str", result.message.content) + content + if chunk.delta.message.tool_calls: + result.message.tool_calls = chunk.delta.message.tool_calls + if chunk.delta.usage: + result.usage.prompt_tokens += chunk.delta.usage.prompt_tokens + result.usage.completion_tokens += chunk.delta.usage.completion_tokens + result.usage.total_tokens += chunk.delta.usage.total_tokens + + result.usage.completion_price = chunk.delta.usage.completion_price + result.usage.prompt_price = chunk.delta.usage.prompt_price + result.usage.total_price = chunk.delta.usage.total_price + result.usage.currency = chunk.delta.usage.currency + result.usage.latency = chunk.delta.usage.latency + + class LLMInvocation(BackwardsInvocation[LLMResultChunk]): @overload def invoke( @@ -84,29 +103,12 @@ def invoke( usage=LLMUsage.empty_usage(), ) - result.message.content = cast("str", result.message.content) - for llm_result in self._backwards_invoke( InvokeType.LLM, LLMResultChunk, data, ): - if isinstance(llm_result.delta.message.content, str): - result.message.content += llm_result.delta.message.content - if len(llm_result.delta.message.tool_calls) > 0: - result.message.tool_calls = llm_result.delta.message.tool_calls - if llm_result.delta.usage: - result.usage.prompt_tokens += llm_result.delta.usage.prompt_tokens - result.usage.completion_tokens += ( - llm_result.delta.usage.completion_tokens - ) - result.usage.total_tokens += llm_result.delta.usage.total_tokens - - result.usage.completion_price = llm_result.delta.usage.completion_price - result.usage.prompt_price = llm_result.delta.usage.prompt_price - result.usage.total_price = llm_result.delta.usage.total_price - result.usage.currency = llm_result.delta.usage.currency - result.usage.latency = llm_result.delta.usage.latency + merge_llm_result_chunk(result, llm_result) return result diff --git a/src/dify_plugin/invocations/model/llm_structured_output.py b/src/dify_plugin/invocations/model/llm_structured_output.py index 9e10f816..99c7a27e 100644 --- a/src/dify_plugin/invocations/model/llm_structured_output.py +++ b/src/dify_plugin/invocations/model/llm_structured_output.py @@ -14,6 +14,7 @@ PromptMessage, PromptMessageTool, ) +from dify_plugin.invocations.model.llm import merge_llm_result_chunk class LLMStructuredOutputInvocation( @@ -99,29 +100,12 @@ def invoke( structured_output=None, ) - result.message.content = cast("str", result.message.content) - for llm_result in self._backwards_invoke( InvokeType.LLMStructuredOutput, LLMResultChunkWithStructuredOutput, data, ): - if isinstance(llm_result.delta.message.content, str): - result.message.content += llm_result.delta.message.content - if len(llm_result.delta.message.tool_calls) > 0: - result.message.tool_calls = llm_result.delta.message.tool_calls - if llm_result.delta.usage: - result.usage.prompt_tokens += llm_result.delta.usage.prompt_tokens - result.usage.completion_tokens += ( - llm_result.delta.usage.completion_tokens - ) - result.usage.total_tokens += llm_result.delta.usage.total_tokens - - result.usage.completion_price = llm_result.delta.usage.completion_price - result.usage.prompt_price = llm_result.delta.usage.prompt_price - result.usage.total_price = llm_result.delta.usage.total_price - result.usage.currency = llm_result.delta.usage.currency - result.usage.latency = llm_result.delta.usage.latency + merge_llm_result_chunk(result, llm_result) # Handle structured output if llm_result.structured_output: diff --git a/src/dify_plugin/invocations/model/tts.py b/src/dify_plugin/invocations/model/tts.py index f98614fc..ca18456b 100644 --- a/src/dify_plugin/invocations/model/tts.py +++ b/src/dify_plugin/invocations/model/tts.py @@ -1,5 +1,6 @@ from binascii import unhexlify from collections.abc import Generator +from contextlib import closing from dify_plugin.core.entities.invocation import InvokeType from dify_plugin.core.runtime import BackwardsInvocation @@ -13,12 +14,14 @@ def invoke( """ Invoke tts """ - for data in self._backwards_invoke( + response = self._backwards_invoke( InvokeType.TTS, TTSResult, { **model_config.model_dump(), "content_text": content_text, }, - ): - yield unhexlify(data.result) + ) + with closing(response): + for data in response: + yield unhexlify(data.result) diff --git a/src/dify_plugin/plugin.py b/src/dify_plugin/plugin.py index 2fda3307..3f0fcfb1 100644 --- a/src/dify_plugin/plugin.py +++ b/src/dify_plugin/plugin.py @@ -2,6 +2,7 @@ import logging import uuid from collections.abc import Generator +from contextlib import closing from typing import Any from pydantic import RootModel @@ -181,16 +182,16 @@ class List(RootModel): ) for file in self.registration.files: - # divide the file into chunks - chunks = [file.data[i : i + 8192] for i in range(0, len(file.data), 8192)] - for sequence, chunk in enumerate(chunks): + for offset in range(0, len(file.data), 8192): tcp_stream.write( InitializeMessage( type=InitializeMessage.Type.ASSET_CHUNK, data=InitializeMessage.AssetChunk( filename=file.filename, - data=base64.b64encode(chunk).decode(), - end=sequence == len(chunks) - 1, + data=base64.b64encode( + file.data[offset : offset + 8192] + ).decode(), + end=offset + 8192 >= len(file.data), ).model_dump(), ).model_dump_json() + "\n\n", @@ -261,6 +262,14 @@ def _register_request_routes(self) -> None: ), ) + self.register_route( + self.plugin_executer.get_tool_runtime_parameters, + lambda data: ( + data.get("type") == PluginInvokeType.Tool.value + and data.get("action") == ToolActions.GetToolRuntimeParameters.value + ), + ) + self.register_route( self.plugin_executer.invoke_agent_strategy, lambda data: ( @@ -569,8 +578,11 @@ def _execute_request( max_invocation_timeout=self.config.MAX_INVOCATION_TIMEOUT, ) response = self.dispatch(session, data) - if response: - if isinstance(response, Generator): + if not response: + return + + if isinstance(response, Generator): + with closing(response): for message in response: if isinstance(message, ToolInvokeMessage) and isinstance( message.message, @@ -622,11 +634,11 @@ def _execute_request( session_id=session_id, data=writer.stream_object(data=message), ) - else: - writer.session_message( - session_id=session_id, - data=writer.stream_object(data=response), - ) + else: + writer.session_message( + session_id=session_id, + data=writer.stream_object(data=response), + ) @staticmethod def _get_remote_install_host_and_port(config: DifyPluginEnv) -> tuple[str, int]: diff --git a/tests/__main__.py b/tests/__main__.py new file mode 100644 index 00000000..998e2511 --- /dev/null +++ b/tests/__main__.py @@ -0,0 +1,8 @@ +from gevent import monkey + +monkey.patch_all(sys=True) + +# pytest plugins can import SSL, so patch before loading pytest. +import pytest # ruff: ignore[module-import-not-at-top-of-file] + +raise SystemExit(pytest.main()) diff --git a/tests/__mock_server/__main__.py b/tests/__mock_server/__main__.py deleted file mode 100644 index d09e43af..00000000 --- a/tests/__mock_server/__main__.py +++ /dev/null @@ -1,13 +0,0 @@ -import sys - -from .openai import openai_server_mock - - -def main() -> None: - sys.stdout.write("OpenAI mock server starting\n") - sys.stdout.flush() - openai_server_mock() - - -if __name__ == "__main__": - main() diff --git a/tests/__mock_server/openai.py b/tests/__mock_server/openai.py deleted file mode 100644 index 85eb9f98..00000000 --- a/tests/__mock_server/openai.py +++ /dev/null @@ -1,47 +0,0 @@ -import json -from collections.abc import Generator - -import flask.cli -from flask import Flask, Response, jsonify, request - -from ..consts.mockserver import OPENAI_MOCK_SERVER_PORT - -flask.cli.show_server_banner = lambda *_args: None - - -def openai_server_mock() -> None: - app = Flask(__name__) - - @app.route("/v1/chat/completions", methods=["POST"]) - def chat_completions() -> Response: - request_body = request.get_json(force=True) - if request_body.get("stream"): - - def stream_response() -> Generator[str, None, None]: - yield "data: " - yield json.dumps({ - "choices": [{"message": {"content": "Hello, world!"}}], - }) - yield "\n\n" - - return Response(stream_response(), mimetype="text/event-stream") - return jsonify({ - "id": "chatcmpl-123", - "object": "chat.completion", - "created": 1715806438, - "model": request_body["model"], - "choices": [ - { - "message": {"role": "assistant", "content": "Hello, world!"}, - "index": 0, - "finish_reason": "stop", - }, - ], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 10, - "total_tokens": 20, - }, - }) - - app.run(host="127.0.0.1", port=OPENAI_MOCK_SERVER_PORT, use_reloader=False) diff --git a/tests/config/test_config.py b/tests/config/test_config.py index 4a42656d..1abe3958 100644 --- a/tests/config/test_config.py +++ b/tests/config/test_config.py @@ -1,51 +1,93 @@ +from pathlib import Path + +import pytest +from pydantic import ValidationError + from dify_plugin.config.config import DifyPluginEnv, InstallMethod -def test_launch() -> None: - """ - Env should works without any parameters and env variables - """ +@pytest.fixture(autouse=True) # ruff: ignore[pytest-fixture-autouse] +def isolated_settings(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + for name in DifyPluginEnv.model_fields: + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv(name.lower(), raising=False) + monkeypatch.chdir(tmp_path) + + +def test_launch_defaults_to_local_install_without_dotenv() -> None: env = DifyPluginEnv() assert InstallMethod.Local == env.INSTALL_METHOD + assert env.REMOTE_INSTALL_URL is None + assert env.REMOTE_INSTALL_KEY is None + assert env.MAX_REQUEST_TIMEOUT > env.MAX_INVOCATION_TIMEOUT > 0 -def test_launch_local_plugin() -> None: - """ - Env should works without any parameters and env variables - """ - env = DifyPluginEnv( - INSTALL_METHOD=InstallMethod.Local, - ) +@pytest.mark.parametrize("method", list(InstallMethod)) +def test_process_environment_selects_install_method_and_parses_limits( + monkeypatch: pytest.MonkeyPatch, method: InstallMethod +) -> None: + monkeypatch.setenv("INSTALL_METHOD", method.value) + monkeypatch.setenv("REMOTE_INSTALL_URL", "127.0.0.1:9001") + monkeypatch.setenv("REMOTE_INSTALL_PORT", "9001") + monkeypatch.setenv("SERVERLESS_PORT", "9002") + monkeypatch.setenv("MAX_REQUEST_TIMEOUT", "7200") + monkeypatch.setenv("MAX_INVOCATION_TIMEOUT", "3700") + monkeypatch.setenv("MAX_WORKER", "3") + monkeypatch.setenv("HEARTBEAT_INTERVAL", "0.25") - assert InstallMethod.Local == env.INSTALL_METHOD + env = DifyPluginEnv() + + assert method == env.INSTALL_METHOD + assert env.REMOTE_INSTALL_URL == "127.0.0.1:9001" + assert env.REMOTE_INSTALL_PORT == 9001 + assert env.SERVERLESS_PORT == 9002 + assert env.MAX_REQUEST_TIMEOUT == 7200 + assert env.MAX_INVOCATION_TIMEOUT == 3700 + assert env.MAX_WORKER == 3 + assert pytest.approx(0.25) == env.HEARTBEAT_INTERVAL -def test_launch_remote_plugin() -> None: - """ - Env should works with remote install url and key - """ - env = DifyPluginEnv( - INSTALL_METHOD=InstallMethod.Remote, - REMOTE_INSTALL_URL="debug.dify.ai:5003", - REMOTE_INSTALL_KEY="19dcf2f3-2856-4fa4-b32b-9ece9b741977", +def test_constructor_overrides_environment_which_overrides_dotenv( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + (tmp_path / ".env").write_text( + "INSTALL_METHOD=remote\n" + "REMOTE_INSTALL_HOST=dotenv.example\n" + "REMOTE_INSTALL_PORT=5004\n" + "MAX_WORKER=2\n" + "PLUGIN_DESCRIPTION=中文\n", + encoding="utf-8", ) + monkeypatch.setenv("INSTALL_METHOD", "serverless") + monkeypatch.setenv("REMOTE_INSTALL_PORT", "9001") + monkeypatch.setenv("MAX_WORKER", "3") - assert InstallMethod.Remote == env.INSTALL_METHOD - assert env.REMOTE_INSTALL_URL == "debug.dify.ai:5003" - assert env.REMOTE_INSTALL_KEY == "19dcf2f3-2856-4fa4-b32b-9ece9b741977" + env = DifyPluginEnv(MAX_WORKER=4) + assert InstallMethod.Serverless == env.INSTALL_METHOD + assert env.REMOTE_INSTALL_HOST == "dotenv.example" + assert env.REMOTE_INSTALL_PORT == 9001 + assert env.MAX_WORKER == 4 + assert "PLUGIN_DESCRIPTION" not in env.model_dump() -def test_launch_serverless_plugin() -> None: - """ - Env should works with serverless install method - """ - env = DifyPluginEnv( - INSTALL_METHOD=InstallMethod.Serverless, - SERVERLESS_HOST="127.0.0.1", - SERVERLESS_PORT=8080, - ) - assert InstallMethod.Serverless == env.INSTALL_METHOD - assert env.SERVERLESS_HOST == "127.0.0.1" - assert env.SERVERLESS_PORT == 8080 +@pytest.mark.parametrize( + ("name", "value"), + [("INSTALL_METHOD", "unknown"), ("MAX_INVOCATION_TIMEOUT", "not-a-number")], +) +def test_invalid_environment_fails_before_startup( + monkeypatch: pytest.MonkeyPatch, name: str, value: str +) -> None: + monkeypatch.setenv(name, value) + + with pytest.raises(ValidationError, match=name): + DifyPluginEnv() + + +def test_settings_cannot_change_after_initialization() -> None: + env = DifyPluginEnv(MAX_WORKER=3) + + with pytest.raises(ValidationError, match="frozen"): + env.MAX_WORKER = 4 + assert env.MAX_WORKER == 3 diff --git a/tests/config/test_integration_config.py b/tests/config/test_integration_config.py new file mode 100644 index 00000000..e4eb693b --- /dev/null +++ b/tests/config/test_integration_config.py @@ -0,0 +1,84 @@ +import sys +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from dify_plugin.config.integration_config import IntegrationConfig, find_dify_cli_path + + +def write_cli(tmp_path: Path, name: str, version: str) -> Path: + cli = tmp_path / name + cli.write_text( + f"#!{sys.executable}\n" + "import sys\n" + "assert sys.argv[1:] == ['version']\n" + f"print({version!r})\n" + ) + cli.chmod(0o700) + return cli + + +@pytest.fixture(autouse=True) # ruff: ignore[pytest-fixture-autouse] +def isolated_cli_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("DIFY_CLI_PATH", raising=False) + monkeypatch.delenv("dify_cli_path", raising=False) + monkeypatch.setenv("PATH", str(tmp_path)) + + +@pytest.mark.parametrize("name", ["dify", "dify-plugin", "dify-plugin-linux-amd64"]) +def test_cli_is_discovered_on_path_and_version_is_checked( + tmp_path: Path, name: str +) -> None: + cli = write_cli(tmp_path, name, "0.1.0") + + assert find_dify_cli_path() == str(cli) + assert IntegrationConfig(_env_file=None).dify_cli_path == str(cli) + + +def test_explicit_cli_overrides_environment_and_path( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + write_cli(tmp_path, "dify", "0.0.1") + environment_cli = write_cli(tmp_path, "environment-cli", "0.1.0") + explicit_cli = write_cli(tmp_path, "explicit-cli", "0.2.0") + monkeypatch.setenv("DIFY_CLI_PATH", str(environment_cli)) + + assert find_dify_cli_path() == str(environment_cli) + assert IntegrationConfig(_env_file=None).dify_cli_path == str(environment_cli) + assert IntegrationConfig( + dify_cli_path=str(explicit_cli), _env_file=None + ).dify_cli_path == str(explicit_cli) + + +def test_missing_cli_has_actionable_configuration_error() -> None: + assert find_dify_cli_path() is None + with pytest.raises(ValidationError, match="dify cli not found"): + IntegrationConfig(_env_file=None) + + +@pytest.mark.parametrize("version", ["0.1.0\n", "0.1.1.dev1", "1.0.0"]) +def test_cli_accepts_supported_versions(tmp_path: Path, version: str) -> None: + cli = write_cli(tmp_path, "custom-cli", version) + + assert IntegrationConfig( + dify_cli_path=str(cli), _env_file=None + ).dify_cli_path == str(cli) + + +@pytest.mark.parametrize( + ("version", "message"), + [ + ("0.0.9", "to support plugin run"), + ("0.1.0rc1", "to support plugin run"), + ("unparseable version", "version is not valid"), + ], +) +def test_cli_rejects_unsupported_or_invalid_version( + tmp_path: Path, version: str, message: str +) -> None: + cli = write_cli(tmp_path, "custom-cli", version) + + with pytest.raises(ValidationError, match=message): + IntegrationConfig(dify_cli_path=str(cli), _env_file=None) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..ebe56e35 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,90 @@ +import os +import re +import ssl +from collections.abc import Generator +from contextlib import ExitStack +from threading import Event + +import pytest +import trustme +from werkzeug.wrappers import Request as HTTPRequest +from werkzeug.wrappers import Response as HTTPResponse + +from tests.http_server import ( + HTTPServer, + HTTPServerFactory, + Request, + Response, + TestHTTPServer, +) + + +@pytest.fixture(autouse=True) # ruff: ignore[pytest-fixture-autouse] +def isolated_http_environment(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY"): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv(name.lower(), raising=False) + for name in ( + "REQUESTS_CA_BUNDLE", + "CURL_CA_BUNDLE", + "SSL_CERT_FILE", + "SSL_CERT_DIR", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("NETRC", os.devnull) + + +@pytest.fixture(scope="session") +def tls_ca() -> trustme.CA: + return trustme.CA() + + +@pytest.fixture +def http_server(tls_ca: trustme.CA) -> Generator[HTTPServerFactory]: + stopped = Event() + servers: list[HTTPServer] = [] + with ExitStack() as stack: + + def start(responses: list[Response], *, tls: bool = False) -> TestHTTPServer: + context = None + if tls: + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + tls_ca.issue_cert("localhost").configure_cert(context) + server = stack.enter_context( + HTTPServer(host="127.0.0.1", threaded=True, ssl_context=context) + ) + servers.append(server) + captured = TestHTTPServer(server.url_for("/").rstrip("/")) + for response in responses: + + def respond( + request: HTTPRequest, response: Response = response + ) -> HTTPResponse: + captured.requests.append( + Request( + request.method, + request.environ["RAW_URI"], + request.headers, + request.get_data(), + ) + ) + headers = dict(response.headers) + if response.chunks is None: + headers.setdefault("Content-Length", str(len(response.body))) + return HTTPResponse( + response.stream(stopped), + status=response.status, + headers=headers, + ) + + server.expect_ordered_request(re.compile(r".*")).respond_with_handler( + respond + ) + return captured + + try: + yield start + for server in servers: + server.check() + finally: + stopped.set() diff --git a/tests/consts/mockserver.py b/tests/consts/mockserver.py deleted file mode 100644 index 19da571e..00000000 --- a/tests/consts/mockserver.py +++ /dev/null @@ -1 +0,0 @@ -OPENAI_MOCK_SERVER_PORT = 11451 diff --git a/tests/core/__init__.py b/tests/core/__init__.py index 8b137891..e69de29b 100644 --- a/tests/core/__init__.py +++ b/tests/core/__init__.py @@ -1 +0,0 @@ - diff --git a/tests/core/plugin_fixtures.py b/tests/core/plugin_fixtures.py new file mode 100644 index 00000000..b35b6a3a --- /dev/null +++ b/tests/core/plugin_fixtures.py @@ -0,0 +1,91 @@ +import json +import shutil +import sys +from collections.abc import Generator +from pathlib import Path + +import pytest + +from dify_plugin import plugin as plugin_module +from dify_plugin.config.config import DifyPluginEnv +from dify_plugin.core.entities.plugin.io import PluginInStream +from dify_plugin.core.server.io_server import RequestReader, ResponseWriter +from dify_plugin.plugin import Plugin + +CONTRACT_SOURCE = Path(__file__).parents[1] / "fixtures/plugins/contract_plugin" + + +class MemoryReader(RequestReader): + def _read_stream(self) -> Generator[PluginInStream, None, None]: + yield from () + + +class CapturingWriter(ResponseWriter): + def __init__(self) -> None: + self.output: list[str] = [] + self.done_count = 0 + + def write(self, data: str) -> None: + self.output.append(data) + + def done(self) -> None: + self.done_count += 1 + + def messages(self) -> list[dict]: + return [json.loads(line) for line in "".join(self.output).splitlines() if line] + + +@pytest.fixture +def contract_source( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> Generator[Path, None, None]: + source = tmp_path / "plugin" + shutil.copytree( + CONTRACT_SOURCE, source, ignore=shutil.ignore_patterns("__pycache__") + ) + monkeypatch.chdir(source) + monkeypatch.syspath_prepend(str(source)) + packages = {p.name for p in source.iterdir() if p.is_dir() and p.name != "_assets"} + original_modules = { + name: module + for name, module in sys.modules.copy().items() + if name.partition(".")[0] in packages + } + for name in original_modules: + del sys.modules[name] + try: + yield source + finally: + for name in tuple(sys.modules): + if name.partition(".")[0] in packages: + del sys.modules[name] + sys.modules.update(original_modules) + + +@pytest.fixture +def contract_plugin( + contract_source: Path, monkeypatch: pytest.MonkeyPatch +) -> Generator[Plugin, None, None]: + del contract_source + monkeypatch.setattr(plugin_module, "StdioRequestReader", MemoryReader) + monkeypatch.setattr(plugin_module, "StdioResponseWriter", CapturingWriter) + plugin = Plugin(DifyPluginEnv(INSTALL_METHOD="local")) + try: + yield plugin + finally: + plugin.close() + plugin.executer.shutdown(wait=True) + + +def invoke(plugin: Plugin, request: dict, **context: object) -> list[dict]: + writer = CapturingWriter() + plugin._execute_request_in_thread( + "session-1", request, plugin.request_reader, writer, **context + ) + messages = writer.messages() + assert writer.done_count == 1 + assert all(message["session_id"] == "session-1" for message in messages) + assert all(message["event"] == "session" for message in messages) + assert messages[-1]["data"] == {"type": "end", "data": {}} + assert sum(message["data"]["type"] == "end" for message in messages) == 1 + return [message["data"] for message in messages[:-1]] diff --git a/tests/core/test_class_loader.py b/tests/core/test_class_loader.py new file mode 100644 index 00000000..3c1b607f --- /dev/null +++ b/tests/core/test_class_loader.py @@ -0,0 +1,188 @@ +import sys +from abc import ABC +from collections.abc import Generator +from pathlib import Path +from types import ModuleType +from uuid import uuid4 + +import pytest + +from dify_plugin.core.utils.class_loader import ( + get_subclasses_from_module, + import_module_from_source, + load_multi_subclasses_from_source, + load_single_subclass_from_source, +) + + +@pytest.fixture +def module_name() -> Generator[str, None, None]: + name = f"test_plugin_{uuid4().hex}" + try: + yield name + finally: + sys.modules.pop(name, None) + + +def test_imports_source_and_registers_module(tmp_path: Path, module_name: str) -> None: + source = tmp_path / "plugin.py" + source.write_text("VALUE = 'loaded'\n", encoding="utf-8") + + module = import_module_from_source( + module_name=module_name, py_file_path=str(source) + ) + + assert module.VALUE == "loaded" + assert sys.modules[module_name] is module + + +def test_lazy_source_runs_when_first_accessed(tmp_path: Path, module_name: str) -> None: + source = tmp_path / "plugin.py" + source.write_text( + "from pathlib import Path\n" + "Path(__file__).with_suffix('.loaded').write_text('ready', encoding='utf-8')\n" + "VALUE = 42\n", + encoding="utf-8", + ) + + module = import_module_from_source( + module_name=module_name, py_file_path=str(source), use_lazy_loader=True + ) + + assert not source.with_suffix(".loaded").exists() + assert module.VALUE == 42 + assert source.with_suffix(".loaded").read_text(encoding="utf-8") == "ready" + + +def test_module_discoverable_on_python_path_loads( + tmp_path: Path, module_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + source = tmp_path / f"{module_name}.py" + source.write_text("VALUE = 42\n", encoding="utf-8") + monkeypatch.syspath_prepend(str(tmp_path)) + + module = import_module_from_source( + module_name=module_name, py_file_path=str(source) + ) + + assert module.VALUE == 42 + + +def test_module_without_source_loader_reports_path( + tmp_path: Path, module_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + namespace = tmp_path / module_name + namespace.mkdir() + monkeypatch.syspath_prepend(str(tmp_path)) + + with pytest.raises(Exception, match=f"Failed to load module {module_name}"): + import_module_from_source(module_name=module_name, py_file_path=str(namespace)) + + +def test_unsupported_source_path_reports_module( + tmp_path: Path, module_name: str +) -> None: + source = tmp_path / "plugin.txt" + source.write_text("VALUE = 42\n", encoding="utf-8") + + with pytest.raises(Exception, match=f"Failed to load module {module_name}"): + import_module_from_source(module_name=module_name, py_file_path=str(source)) + + +@pytest.mark.parametrize( + ("source_text", "error_type", "message"), + [ + ("class Broken(\n", SyntaxError, "never closed"), + ("raise RuntimeError('plugin failed')\n", RuntimeError, "plugin failed"), + ], +) +def test_source_errors_reach_caller( + tmp_path: Path, + module_name: str, + source_text: str, + error_type: type[Exception], + message: str, +) -> None: + source = tmp_path / "plugin.py" + source.write_text(source_text, encoding="utf-8") + + with pytest.raises(error_type, match=message): + import_module_from_source(module_name=module_name, py_file_path=str(source)) + + +def test_subclass_discovery_excludes_parent_instances_and_unrelated_types() -> None: + class Parent: + pass + + class Child(Parent): + pass + + class Grandchild(Child): + pass + + module = ModuleType("plugin") + module.Parent = Parent + module.Child = Child + module.Grandchild = Grandchild + module.Unrelated = str + module.instance = Child() + + assert get_subclasses_from_module(module, Parent) == [Child, Grandchild] + + +@pytest.mark.parametrize("lazy", [False, True]) +def test_loads_multiple_plugin_classes( + tmp_path: Path, module_name: str, lazy: bool +) -> None: + source = tmp_path / "plugin.py" + source.write_text( + "from abc import ABC\nclass Alpha(ABC): pass\nclass Beta(ABC): pass\n", + encoding="utf-8", + ) + + classes = load_multi_subclasses_from_source( + module_name=module_name, + script_path=str(source), + parent_type=ABC, + use_lazy_loader=lazy, + ) + + assert [cls.__name__ for cls in classes] == ["Alpha", "Beta"] + assert all(issubclass(cls, ABC) for cls in classes) + + +@pytest.mark.parametrize( + ("definitions", "error"), + [ + ("", "Missing subclass of ABC"), + ( + "class Alpha(ABC): pass\nclass Beta(ABC): pass\n", + "Multiple subclasses of ABC", + ), + ], +) +def test_single_plugin_class_rejects_missing_or_ambiguous_implementations( + tmp_path: Path, module_name: str, definitions: str, error: str +) -> None: + source = tmp_path / "plugin.py" + source.write_text(f"from abc import ABC\n{definitions}", encoding="utf-8") + + with pytest.raises(Exception, match=error): + load_single_subclass_from_source( + module_name=module_name, script_path=str(source), parent_type=ABC + ) + + +def test_single_plugin_class_can_be_instantiated( + tmp_path: Path, module_name: str +) -> None: + source = tmp_path / "plugin.py" + source.write_text( + "from abc import ABC\nclass Plugin(ABC): VALUE = 42\n", encoding="utf-8" + ) + + implementation = load_single_subclass_from_source( + module_name=module_name, script_path=str(source), parent_type=ABC + ) + + assert implementation().VALUE == 42 diff --git a/tests/core/test_endpoint_lifecycle.py b/tests/core/test_endpoint_lifecycle.py new file mode 100644 index 00000000..6ac67875 --- /dev/null +++ b/tests/core/test_endpoint_lifecycle.py @@ -0,0 +1,126 @@ +from collections.abc import Generator +from pathlib import Path + +import pytest +from werkzeug import Response + +from dify_plugin.plugin import Plugin + +from .plugin_fixtures import CapturingWriter, contract_plugin, contract_source, invoke + +__all__ = ["contract_plugin", "contract_source"] + +ENDPOINT_REQUEST = { + "type": "endpoint", + "action": "invoke_endpoint", + "raw_http_request": b"POST /echo/item HTTP/1.1\r\nHost: localhost\r\n\r\n".hex(), + "settings": {}, +} + + +@pytest.mark.parametrize("outcome", ["complete", "provider-error", "disconnect"]) +def test_endpoint_closes_stream_and_response_when_request_ends( + contract_plugin: Plugin, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + outcome: str, +) -> None: + resource = (tmp_path / "response.bin").open("w+b") + resource.write(b"firstsecond") + resource.seek(0) + closed: list[str] = [] + + def chunks() -> Generator[bytes]: + try: + yield resource.read(5) + if outcome == "provider-error": + msg = "upstream response interrupted" + raise OSError(msg) + yield resource.read() + finally: + closed.append("iterator") + + response = Response(chunks(), content_type="application/octet-stream") + response.call_on_close(resource.close) + response.call_on_close(lambda: closed.append("response")) + endpoint = next(contract_plugin.registration.endpoints.iter_rules()).endpoint + monkeypatch.setattr(endpoint, "_invoke", lambda *_args: response) + writer = CapturingWriter() + + if outcome == "disconnect": + write = writer.write + + def disconnect(data: str) -> None: + if len(writer.output) == 2: + msg = "daemon disconnected after first body chunk" + raise BrokenPipeError(msg) + write(data) + + monkeypatch.setattr(writer, "write", disconnect) + + try: + if outcome == "disconnect": + with pytest.raises(BrokenPipeError, match="daemon disconnected"): + contract_plugin._execute_request_in_thread( + "session-1", + ENDPOINT_REQUEST, + contract_plugin.request_reader, + writer, + ) + else: + contract_plugin._execute_request_in_thread( + "session-1", ENDPOINT_REQUEST, contract_plugin.request_reader, writer + ) + + messages = [message["data"] for message in writer.messages()] + assert messages[0]["data"]["status"] == 200 + assert messages[1]["data"]["result"] == b"first".hex() + assert writer.done_count == 1 + if outcome == "complete": + assert [message["type"] for message in messages] == [ + "stream", + "stream", + "stream", + "end", + ] + assert messages[2]["data"]["result"] == b"second".hex() + elif outcome == "provider-error": + assert [message["type"] for message in messages] == [ + "stream", + "stream", + "error", + "end", + ] + assert messages[2]["data"]["error_type"] == "OSError" + assert messages[2]["data"]["message"] == "upstream response interrupted" + else: + assert [message["type"] for message in messages] == ["stream", "stream"] + + # Werkzeug Response.close owns both iterator cleanup and close callbacks. + assert closed == ["iterator", "response"] + assert resource.closed + finally: + response.close() + resource.close() + + +def test_endpoint_closes_file_iterable_after_sending_its_body( + contract_plugin: Plugin, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "response.txt" + path.write_bytes(b"first\nsecond\n") + resource = path.open("rb") + response = Response(resource, content_type="text/plain") + endpoint = next(contract_plugin.registration.endpoints.iter_rules()).endpoint + monkeypatch.setattr(endpoint, "_invoke", lambda *_args: response) + + try: + messages = invoke(contract_plugin, ENDPOINT_REQUEST) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert messages[0]["data"]["status"] == 200 + assert bytes.fromhex(messages[0]["data"]["result"]) == path.read_bytes() + assert resource.closed + finally: + response.close() diff --git a/tests/core/test_invoke_speech_to_text.py b/tests/core/test_invoke_speech_to_text.py index 2812c1a8..c0777f34 100644 --- a/tests/core/test_invoke_speech_to_text.py +++ b/tests/core/test_invoke_speech_to_text.py @@ -30,6 +30,59 @@ ), (b"\x1a\x45\xdf\xa3\x88\x42\x82\x40\x04webm", ".webm"), (b"\x1a\x45\xdf\xa3\x88\x42\x82\x85webm\0", ".webm"), + pytest.param( + bytes.fromhex("1a45dfa301000000000000074282847765626d"), + ".webm", + id="webm-eight-byte-header-size", + ), + pytest.param( + bytes.fromhex("1a45dfa389ec804282847765626d"), + ".webm", + id="webm-after-empty-void", + ), + pytest.param(bytes.fromhex("1a45dfa3"), ".mp3", id="ebml-missing-header-size"), + pytest.param( + bytes.fromhex("1a45dfa340"), ".mp3", id="ebml-truncated-header-size" + ), + pytest.param( + bytes.fromhex("1a45dfa300"), ".mp3", id="ebml-invalid-size-marker" + ), + pytest.param( + bytes.fromhex("1a45dfa3ff4282847765626d"), + ".mp3", + id="ebml-unknown-header-size", + ), + pytest.param( + bytes.fromhex("1a45dfa388"), ".mp3", id="ebml-truncated-header-data" + ), + pytest.param(bytes.fromhex("1a45dfa380"), ".mp3", id="ebml-empty-header"), + pytest.param( + bytes.fromhex("1a45dfa38100"), ".mp3", id="ebml-invalid-element-marker" + ), + pytest.param( + bytes.fromhex("1a45dfa38142"), ".mp3", id="ebml-truncated-element-id" + ), + pytest.param( + bytes.fromhex("1a45dfa3824282"), ".mp3", id="ebml-missing-element-size" + ), + pytest.param( + bytes.fromhex("1a45dfa3874282ff7765626d"), + ".mp3", + id="ebml-unknown-doctype-size", + ), + pytest.param( + bytes.fromhex("1a45dfa3864282847765626d"), + ".mp3", + id="ebml-doctype-exceeds-parent", + ), + pytest.param( + bytes.fromhex("1a45dfa38a4282877765626d"), + ".mp3", + id="ebml-doctype-exceeds-file", + ), + pytest.param( + bytes.fromhex("1a45dfa382ec80"), ".mp3", id="ebml-without-doctype" + ), (b"RIFF\x24\x08\x00\x00AVI payload", ".mp3"), (b"\x00\x00\x00\x18ftypavifpayload", ".mp3"), (b"\x1a\x45\xdf\xa3\x9f\x42\x82\x88matroskapayload", ".mp3"), @@ -49,6 +102,8 @@ def test_invoke_speech_to_text_preserves_audio_format( audio: bytes, expected_suffix: str, ) -> None: + files: list[IO[bytes]] = [] + def invoke( model_name: str, credentials: dict[str, object], @@ -60,6 +115,7 @@ def invoke( assert pathlib.Path(audio_file.name).suffix == expected_suffix assert audio_file.tell() == 0 assert audio_file.read() == audio + files.append(audio_file) return "transcript" model = MagicMock(spec=Speech2TextModel) @@ -79,3 +135,6 @@ def invoke( assert executor.invoke_speech_to_text(Session.empty_session(), request) == { "result": "transcript", } + assert len(files) == 1 + assert files[0].closed + assert not pathlib.Path(files[0].name).exists() diff --git a/tests/core/test_plugin_initialization.py b/tests/core/test_plugin_initialization.py new file mode 100644 index 00000000..f9a8c02b --- /dev/null +++ b/tests/core/test_plugin_initialization.py @@ -0,0 +1,223 @@ +import base64 +from collections.abc import Callable +from pathlib import Path + +import pytest + +from dify_plugin import plugin as plugin_module +from dify_plugin.config.config import DifyPluginEnv +from dify_plugin.core.entities.plugin.setup import PluginAsset +from dify_plugin.plugin import Plugin + +from .plugin_fixtures import ( + CapturingWriter, + MemoryReader, + contract_plugin, + contract_source, +) + +__all__ = ["contract_plugin", "contract_source"] + + +def test_local_initialization_writes_manifest_first(contract_plugin: Plugin) -> None: + writer = contract_plugin.default_response_writer + assert isinstance(writer, CapturingWriter) + messages = writer.messages() + assert len(messages) == 1 + assert messages[0]["name"] == "contract_plugin" + assert messages[0]["plugins"]["tools"] == ["provider/contract.yaml"] + assert messages[0]["plugins"]["triggers"] == ["triggers/events.yaml"] + assert writer.output[0].endswith("\n\n") + + +def test_remote_initialization_sends_all_declarations_before_assets( + contract_plugin: Plugin, +) -> None: + stream = CapturingWriter() + contract_plugin._initialize_tcp_stream(stream) + messages = stream.messages() + + assert [message["type"] for message in messages[:7]] == [ + "manifest_declaration", + "tool_declaration", + "model_declaration", + "endpoint_declaration", + "agent_strategy_declaration", + "datasource_declaration", + "trigger_declaration", + ] + assert messages[0]["data"]["name"] == "contract_plugin" + assert messages[1]["data"][0]["identity"]["name"] == "contract" + assert messages[1]["data"][0]["tools"][0]["has_runtime_parameters"] is True + assert messages[2]["data"][0]["provider"] == "contract_models" + assert messages[3]["data"][0]["endpoints"][0]["path"] == "/echo/" + assert messages[4]["data"][0]["identity"]["name"] == "contract_agents" + assert {provider["identity"]["name"] for provider in messages[5]["data"]} == { + "website", + "documents", + "drive", + } + assert messages[6]["data"][0]["identity"]["name"] == "events" + assert all(message["type"] == "asset_chunk" for message in messages[7:-1]) + assert messages[-1] == {"type": "end", "data": {}} + assert all(frame.endswith("\n\n") for frame in stream.output) + + +@pytest.mark.parametrize("size", [0, 8191, 8192, 8193, 16384]) +def test_remote_asset_chunks_preserve_each_file( + contract_plugin: Plugin, size: int +) -> None: + content = b"\x00\xff" * (size // 2) + b"x" * (size % 2) + contract_plugin.registration.files = [ + PluginAsset(filename="test.bin", data=content) + ] + stream = CapturingWriter() + + contract_plugin._initialize_tcp_stream(stream) + + chunks = [ + message["data"] + for message in stream.messages() + if message["type"] == "asset_chunk" + ] + assert len(chunks) == (size + 8191) // 8192 + assert all(chunk["filename"] == "test.bin" for chunk in chunks) + assert all(chunk["data"] for chunk in chunks) + assert [chunk["end"] for chunk in chunks] == ( + [False] * (len(chunks) - 1) + [True] if chunks else [] + ) + assert b"".join(base64.b64decode(chunk["data"]) for chunk in chunks) == content + + +def test_remote_initialization_omits_undeclared_extensions( + contract_plugin: Plugin, +) -> None: + registration = contract_plugin.registration + registration.tools_configuration.clear() + registration.models_configuration.clear() + registration.endpoints_configuration.clear() + registration.agent_strategies_configuration.clear() + registration.datasource_configuration.clear() + registration.triggers_configuration.clear() + registration.files.clear() + stream = CapturingWriter() + + contract_plugin._initialize_tcp_stream(stream) + + assert [message["type"] for message in stream.messages()] == [ + "manifest_declaration", + "end", + ] + + +def test_remote_constructor_launches_transport_with_registration_callback( + contract_source: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + del contract_source + connections = [] + + class RemoteStream(MemoryReader, CapturingWriter): + def __init__( + self, host: str, port: int, key: str, on_connected: Callable + ) -> None: + MemoryReader.__init__(self) + CapturingWriter.__init__(self) + self.on_connected = on_connected + connections.append((host, port, key)) + + def launch(self) -> None: + self.on_connected() + + monkeypatch.setattr(plugin_module, "TCPReaderWriter", RemoteStream) + plugin = Plugin( + DifyPluginEnv( + INSTALL_METHOD="remote", + REMOTE_INSTALL_URL="tcp://127.0.0.1:6500", + REMOTE_INSTALL_KEY="test-key", + ) + ) + try: + assert connections == [("127.0.0.1", 6500, "test-key")] + assert plugin.request_reader is plugin.default_response_writer + messages = plugin.default_response_writer.messages() + assert messages[0]["type"] == "manifest_declaration" + assert messages[-1]["type"] == "end" + finally: + plugin.close() + plugin.executer.shutdown(wait=True) + + +def test_remote_constructor_requires_install_key(contract_source: Path) -> None: + del contract_source + with pytest.raises(ValueError, match="Missing remote install key"): + Plugin(DifyPluginEnv(INSTALL_METHOD="remote", REMOTE_INSTALL_KEY=None)) + + +def test_serverless_constructor_passes_worker_configuration( + contract_source: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + del contract_source + options = {} + launches = [] + + class ServerlessStream(MemoryReader): + def __init__(self, **kwargs: object) -> None: + super().__init__() + options.update(kwargs) + + def launch(self) -> None: + launches.append(True) + + monkeypatch.setattr(plugin_module, "ServerlessRequestReader", ServerlessStream) + plugin = Plugin( + DifyPluginEnv( + INSTALL_METHOD="serverless", + SERVERLESS_HOST="127.0.0.1", + SERVERLESS_PORT=9000, + SERVERLESS_WORKERS=2, + SERVERLESS_THREADS=3, + SERVERLESS_WORKER_CONNECTIONS=20, + MAX_REQUEST_TIMEOUT=30, + ) + ) + try: + assert launches == [True] + assert options == { + "host": "127.0.0.1", + "port": 9000, + "worker_class": "gevent", + "workers": 2, + "threads": 3, + "worker_connections": 20, + "max_single_connection_lifetime": 30, + } + assert plugin.default_response_writer is None + assert len(plugin.routes) == 37 + finally: + plugin.close() + plugin.executer.shutdown(wait=True) + + +@pytest.mark.parametrize( + ("url", "expected"), + [ + (None, ("configured-host", 7000)), + ("host:6500", ("host", 6500)), + ("tcp://host:6500", ("host", 6500)), + ("http://host", ("host", 80)), + ("tcp://[::1]:6500", ("::1", 6500)), + ], +) +def test_remote_address_formats(url: str | None, expected: tuple[str, int]) -> None: + config = DifyPluginEnv( + REMOTE_INSTALL_URL=url, + REMOTE_INSTALL_HOST="configured-host", + REMOTE_INSTALL_PORT=7000, + ) + assert Plugin._get_remote_install_host_and_port(config) == expected + + +@pytest.mark.parametrize("url", ["host", "host:not-a-port"]) +def test_remote_address_rejects_invalid_port_configuration(url: str) -> None: + with pytest.raises(ValueError, match=r"Invalid remote install URL|invalid literal"): + Plugin._get_remote_install_host_and_port(DifyPluginEnv(REMOTE_INSTALL_URL=url)) diff --git a/tests/core/test_plugin_registration.py b/tests/core/test_plugin_registration.py index de0877ee..b603d098 100644 --- a/tests/core/test_plugin_registration.py +++ b/tests/core/test_plugin_registration.py @@ -1,91 +1,146 @@ -from collections import UserList -from unittest.mock import MagicMock +from pathlib import Path import pytest +import yaml -from dify_plugin.core import plugin_registration as registration_module +from dify_plugin.config.config import DifyPluginEnv from dify_plugin.core.plugin_registration import PluginRegistration +from .plugin_fixtures import contract_source -class AppendTrackingList(UserList): - def __init__(self) -> None: - super().__init__() - self.appended: list[object] = [] +__all__ = ["contract_source"] - def append(self, item: object) -> None: - self.appended.append(item) - super().append(item) + +def remove_constructor_export(source: Path) -> None: + dispatch = source / "triggers/dispatch.py" + dispatch.write_text( + dispatch.read_text().replace( + "from triggers.constructor import ContractConstructor\n", "" + ) + ) + + +def test_subscription_constructor_loads_its_declared_source( + contract_source: Path, +) -> None: + remove_constructor_export(contract_source) + provider_file = contract_source / "triggers/events.yaml" + provider = yaml.safe_load(provider_file.read_text()) + provider["subscription_constructor"]["extra"] = { + "python": {"source": "triggers/constructor.py"} + } + provider_file.write_text(yaml.safe_dump(provider)) + + registration = PluginRegistration(DifyPluginEnv()) + + constructor = registration.trigger_factory.get_subscription_constructor_cls( + "events" + ) + assert constructor is not None + assert constructor.__name__ == "ContractConstructor" + assert constructor.__module__ == "triggers.constructor" -def test_configuration_loading_preserves_append_hooks( - monkeypatch: pytest.MonkeyPatch, +def test_subscription_constructor_keeps_provider_source_precedence( + contract_source: Path, ) -> None: - configuration = MagicMock() - configuration.plugins.tools = ["tool.yaml"] - configuration.plugins.models = [] - configuration.plugins.endpoints = [] - configuration.plugins.agent_strategies = [] - configuration.plugins.datasources = [] - configuration.plugins.triggers = [] - - def load_configuration(file_path: str, configuration_type: type) -> object: - del configuration_type - if file_path == "manifest.yaml": - return configuration, {} - return file_path, {} - - monkeypatch.setattr( - registration_module, - "_load_configuration", - load_configuration, + provider_file = contract_source / "triggers/events.yaml" + provider = yaml.safe_load(provider_file.read_text()) + provider["subscription_constructor"]["extra"] = { + "python": {"source": "triggers/unused_missing_constructor.py"} + } + provider_file.write_text(yaml.safe_dump(provider)) + + registration = PluginRegistration(DifyPluginEnv()) + + constructor = registration.trigger_factory.get_subscription_constructor_cls( + "events" ) - registration = object.__new__(PluginRegistration) - registration.tools_configuration = AppendTrackingList() - registration.models_configuration = [] - registration.endpoints_configuration = [] - registration.agent_strategies_configuration = [] - registration.datasource_configuration = [] - registration.triggers_configuration = [] + assert constructor is not None + assert constructor.__module__ == "triggers.constructor" - registration._load_plugin_configuration() - assert registration.tools_configuration.appended == ["tool.yaml"] +def test_subscription_constructor_requires_one_implementation( + contract_source: Path, +) -> None: + remove_constructor_export(contract_source) + + with pytest.raises( + ValueError, match="configuration declared but no implementation" + ): + PluginRegistration(DifyPluginEnv()) + + +def test_subscription_constructor_rejects_multiple_implementations( + contract_source: Path, +) -> None: + dispatch = contract_source / "triggers/dispatch.py" + dispatch.write_text( + dispatch.read_text() + + "\nclass DuplicateConstructor(ContractConstructor):\n pass\n" + ) + + with pytest.raises(ValueError, match="Multiple TriggerSubscriptionConstructor"): + PluginRegistration(DifyPluginEnv()) -def test_configuration_list_is_bound_after_provider_loading( - monkeypatch: pytest.MonkeyPatch, +def test_trigger_without_subscription_configuration_is_supported( + contract_source: Path, ) -> None: - configuration = MagicMock() - configuration.plugins.tools = ["tool.yaml"] - configuration.plugins.models = [] - configuration.plugins.endpoints = [] - configuration.plugins.agent_strategies = [] - configuration.plugins.datasources = [] - configuration.plugins.triggers = [] - registration = object.__new__(PluginRegistration) - old_list = AppendTrackingList() - new_list = AppendTrackingList() - registration.tools_configuration = old_list - registration.models_configuration = [] - registration.endpoints_configuration = [] - registration.agent_strategies_configuration = [] - registration.datasource_configuration = [] - registration.triggers_configuration = [] - - def load_configuration(file_path: str, configuration_type: type) -> object: - del configuration_type - if file_path == "manifest.yaml": - return configuration, {} - registration.tools_configuration = new_list - return file_path, {} - - monkeypatch.setattr( - registration_module, - "_load_configuration", - load_configuration, + remove_constructor_export(contract_source) + provider_file = contract_source / "triggers/events.yaml" + provider = yaml.safe_load(provider_file.read_text()) + provider.pop("subscription_constructor") + provider_file.write_text(yaml.safe_dump(provider)) + + registration = PluginRegistration(DifyPluginEnv()) + + assert ( + registration.trigger_factory.get_subscription_constructor_cls("events") is None ) + assert ( + registration.trigger_factory.get_configuration("events").events[0].identity.name + == "echo" + ) + + +@pytest.mark.parametrize("missing_file", ["manifest.yaml", "provider/contract.yaml"]) +def test_missing_configuration_has_context_and_cause( + contract_source: Path, missing_file: str +) -> None: + (contract_source / missing_file).unlink() + + with pytest.raises(ValueError, match="Error loading plugin configuration") as error: + PluginRegistration(DifyPluginEnv()) + + assert error.value.__cause__ is not None + assert missing_file in str(error.value) + + +def test_invalid_manifest_does_not_start_loading_implementations( + contract_source: Path, +) -> None: + manifest = contract_source / "manifest.yaml" + manifest.write_text("version: invalid\n") + + with pytest.raises(ValueError, match="Error loading plugin configuration") as error: + PluginRegistration(DifyPluginEnv()) + + assert "validation error" in str(error.value) + + +def test_assets_load_binary_files_without_directories(contract_source: Path) -> None: + assets = contract_source / "_assets" + (assets / "empty.bin").touch() + (assets / "binary.bin").write_bytes(bytes(range(256))) + nested = assets / "nested" + nested.mkdir() + (nested / "ignored.txt").write_text("ignored") - registration._load_plugin_configuration() + registration = PluginRegistration(DifyPluginEnv()) - assert old_list.appended == [] - assert new_list.appended == ["tool.yaml"] + files = {asset.filename: asset.data for asset in registration.files} + assert files["empty.bin"] == b"" + assert files["binary.bin"] == bytes(range(256)) + assert "nested" not in files + assert "ignored.txt" not in files diff --git a/tests/core/test_plugin_routes.py b/tests/core/test_plugin_routes.py new file mode 100644 index 00000000..13b04334 --- /dev/null +++ b/tests/core/test_plugin_routes.py @@ -0,0 +1,1468 @@ +import base64 +import hmac +from collections.abc import Generator +from copy import deepcopy +from inspect import signature +from io import BytesIO +from typing import Any + +import pytest +from pydantic import TypeAdapter +from pytest_mock import MockerFixture +from werkzeug import Request, Response +from werkzeug.wsgi import FileWrapper + +from dify_plugin.core.entities.plugin.request import ( + AgentActions, + DatasourceActions, + DynamicParameterActions, + EndpointActions, + ModelActions, + OAuthActions, + PluginInvokeType, + ToolActions, + TriggerActions, +) +from dify_plugin.entities import I18nObject, ParameterOption +from dify_plugin.entities.agent import AgentInvokeMessage +from dify_plugin.entities.datasource import DatasourceMessage +from dify_plugin.entities.model import ModelType +from dify_plugin.entities.tool import ToolInvokeMessage +from dify_plugin.entities.trigger import EventDispatch, Subscription +from dify_plugin.interfaces.tool import Tool +from dify_plugin.invocations.file import File +from dify_plugin.plugin import Plugin + +from .plugin_fixtures import contract_plugin, contract_source, invoke + +__all__ = ["contract_plugin", "contract_source"] + +ACTIONS = { + PluginInvokeType.Tool: ToolActions, + PluginInvokeType.Model: ModelActions, + PluginInvokeType.Agent: AgentActions, + PluginInvokeType.Endpoint: EndpointActions, + PluginInvokeType.OAuth: OAuthActions, + PluginInvokeType.Datasource: DatasourceActions, + PluginInvokeType.Trigger: TriggerActions, + PluginInvokeType.DynamicParameter: DynamicParameterActions, +} + + +CREDENTIALS = {"api_key": "valid"} +SUBSCRIPTION = { + "endpoint": "https://plugin.invalid/events", + "expires_at": 100, + "parameters": {"event": "echo"}, + "properties": {"external_id": "subscription"}, +} + + +def raw_request(path: str, body: bytes = b"", method: str = "POST") -> str: + return ( + f"{method} {path} HTTP/1.1\r\nHost: localhost\r\n" + f"Content-Type: application/json\r\nContent-Length: {len(body)}\r\n\r\n" + ).encode().hex() + body.hex() + + +# Wire names and request fields follow daemon access_types/access.go and requests/. +CASES = [ + ( + "tool", + "invoke_tool", + { + "provider": "contract", + "tool": "echo", + "tool_parameters": {"message": "hello"}, + "credential_type": "oauth2", + }, + [ + {"type": "text", "message": {"text": "hello"}}, + {"type": "json", "message": {"json_object": {"echo": "hello"}}}, + ], + ), + ("tool", "validate_tool_credentials", {"provider": "contract"}, [{"result": True}]), + ( + "tool", + "get_tool_runtime_parameters", + {"provider": "contract", "tool": "echo"}, + [ + { + "parameters": [ + {"name": "message", "type": "string", "required": True}, + {"name": "wait_for_peer", "type": "boolean", "required": False}, + ] + } + ], + ), + ( + "agent_strategy", + "invoke_agent_strategy", + { + "agent_strategy_provider": "contract_agents", + "agent_strategy": "echo", + "agent_strategy_params": {"message": "hello"}, + }, + [{"type": "text", "message": {"text": "hello"}}], + ), + ( + "endpoint", + "invoke_endpoint", + { + "raw_http_request": raw_request("/echo/item", b"hello"), + "settings": {"label": "setting"}, + }, + [{"status": 201, "result": b"hello".hex()}], + ), + ( + "oauth", + "get_authorization_url", + { + "provider": "events", + "redirect_uri": "https://plugin.invalid/callback", + "system_credentials": {"client_id": "client"}, + }, + [ + { + "authorization_url": "https://provider.invalid/authorize?redirect_uri=https%3A%2F%2Fplugin.invalid%2Fcallback&client_id=client" + } + ], + ), + ( + "oauth", + "get_credentials", + { + "provider": "events", + "redirect_uri": "https://plugin.invalid/callback", + "system_credentials": {"client_id": "client"}, + "raw_http_request": raw_request("/callback?code=code", method="GET"), + }, + [ + { + "credentials": { + "access_token": "token-code", + "refresh_token": "refresh", + }, + "expires_at": -1, + "metadata": {}, + } + ], + ), + ( + "oauth", + "refresh_credentials", + { + "provider": "events", + "redirect_uri": "https://plugin.invalid/callback", + "system_credentials": {"client_id": "client"}, + "credentials": {"access_token": "old", "refresh_token": "refresh"}, + }, + [ + { + "credentials": { + "access_token": "refreshed", + "refresh_token": "refresh", + }, + "expires_at": -1, + } + ], + ), + ( + "datasource", + "validate_datasource_credentials", + {"provider": "website"}, + [{"result": True}], + ), + ( + "datasource", + "invoke_website_datasource_get_crawl", + { + "provider": "website", + "datasource": "echo", + "datasource_parameters": {"url": "https://example.invalid/page"}, + }, + [ + {"result": {"status": "processing", "completed": 0}}, + { + "result": { + "status": "completed", + "completed": 1, + "web_info_list": [ + { + "source_url": "https://example.invalid/page", + "content": "page content", + } + ], + }, + }, + ], + ), + ( + "datasource", + "invoke_online_document_datasource_get_pages", + { + "provider": "documents", + "datasource": "echo", + "datasource_parameters": {"workspace": "selected", "limit": 0}, + }, + [ + { + "result": [ + { + "workspace_id": "workspace", + "pages": [{"page_id": "page", "page_name": "Example"}], + } + ] + } + ], + ), + ( + "datasource", + "invoke_online_document_datasource_get_page_content", + { + "provider": "documents", + "datasource": "echo", + "page": {"workspace_id": "workspace", "page_id": "page", "type": "page"}, + }, + [ + { + "type": "variable", + "message": { + "variable_name": "content", + "variable_value": "page content", + }, + }, + { + "type": "variable", + "message": {"variable_name": "page_id", "variable_value": "page"}, + }, + { + "type": "variable", + "message": { + "variable_name": "workspace_id", + "variable_value": "workspace", + }, + }, + ], + ), + ( + "datasource", + "invoke_online_drive_browse_files", + { + "provider": "drive", + "datasource": "echo", + "request": { + "bucket": "bucket", + "prefix": "folder/", + "max_keys": 20, + "next_page_parameters": {}, + }, + }, + [ + { + "result": [ + { + "bucket": "bucket", + "files": [{"id": "file", "name": "example.txt"}], + "is_truncated": True, + "next_page_parameters": {"cursor": "next"}, + } + ] + } + ], + ), + ( + "datasource", + "invoke_online_drive_download_file", + { + "provider": "drive", + "datasource": "echo", + "request": {"bucket": "bucket", "id": "file"}, + }, + [ + { + "type": "blob", + "message": {"blob": "ZGF0YQ=="}, + "meta": {"mime_type": "text/plain", "file_name": "example.txt"}, + }, + ], + ), + ( + "trigger", + "invoke_trigger_event", + { + "provider": "events", + "event": "echo", + "subscription": SUBSCRIPTION, + "raw_http_request": raw_request("/events", b'{"message":"hello"}'), + "parameters": {"field": "message"}, + "payload": {"message": "hello"}, + "credential_type": "oauth2", + }, + [ + { + "variables": { + "payload": {"message": "hello"}, + "parameters": {"field": "message"}, + "body": '{"message":"hello"}', + }, + "cancelled": False, + } + ], + ), + ( + "trigger", + "validate_trigger_credentials", + {"provider": "events"}, + [{"result": True}], + ), + ( + "trigger", + "dispatch_trigger_event", + { + "provider": "events", + "subscription": SUBSCRIPTION, + "raw_http_request": raw_request( + "/events", b'{"message":"hello","sender":"external-user"}' + ), + "credential_type": "oauth2", + }, + [ + { + "events": ["echo"], + "payload": {"message": "hello", "sender": "external-user"}, + "user_id": "external-user", + } + ], + ), + ( + "trigger", + "subscribe_trigger", + { + "provider": "events", + "endpoint": "https://plugin.invalid/events", + "parameters": {"event": "echo"}, + "credential_type": "api-key", + }, + [ + { + "subscription": { + **SUBSCRIPTION, + "properties": { + "external_id": "subscription", + "credential_type": "api-key", + }, + } + } + ], + ), + ( + "trigger", + "unsubscribe_trigger", + { + "provider": "events", + "subscription": SUBSCRIPTION, + "credential_type": "api-key", + }, + [{"subscription": {"success": True, "message": "subscription"}}], + ), + ( + "trigger", + "refresh_trigger", + { + "provider": "events", + "subscription": SUBSCRIPTION, + "credential_type": "api-key", + }, + [{"subscription": {**SUBSCRIPTION, "expires_at": 200}}], + ), + ( + "dynamic_parameter", + "fetch_parameter_options", + { + "provider": "events", + "provider_action": "provider", + "parameter": "event", + "credential_type": "api-key", + }, + [{"options": [{"value": "echo", "label": {"en_US": "event"}}]}], + ), +] + +MODEL_CASES = [ + ( + "validate_provider_credentials", + "llm", + {}, + [{"result": True, "credentials": CREDENTIALS}], + ), + ( + "validate_model_credentials", + "llm", + {}, + [{"result": True, "credentials": CREDENTIALS}], + ), + ( + "invoke_llm", + "llm", + { + "prompt_messages": [{"role": "user", "content": "hello"}], + "model_parameters": {"temperature": 0.5}, + "tools": [ + { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + } + ], + "stop": ["END"], + "stream": True, + }, + [{"model": "echo", "delta": {"message": {"content": "hello"}}}], + ), + ( + "start_polling", + "llm", + { + "prompt_messages": [{"role": "user", "content": "hello"}], + "model_parameters": {"temperature": 0.5}, + "tools": [ + { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + } + ], + "stop": ["END"], + "stream": False, + "json_schema": {"type": "object"}, + }, + [ + { + "status": "running", + "plugin_state": { + "model": "echo", + "content": "hello", + "json_schema": {"type": "object"}, + }, + } + ], + ), + ( + "check_polling", + "llm", + {"plugin_state": {"content": "hello"}}, + [ + { + "status": "succeeded", + "result": {"model": "echo", "message": {"content": "hello"}}, + } + ], + ), + ( + "get_llm_num_tokens", + "llm", + {"prompt_messages": [{"role": "user", "content": "hello"}], "tools": []}, + [{"num_tokens": 5}], + ), + ( + "invoke_text_embedding", + "text-embedding", + {"texts": ["hello"], "input_type": "query"}, + [{"model": "echo", "embeddings": [[5.0]]}], + ), + ( + "invoke_multimodal_embedding", + "text-embedding", + { + "documents": [{"content": "hello", "content_type": "text"}], + "input_type": "query", + }, + [{"model": "echo", "embeddings": [[5.0]]}], + ), + ( + "get_text_embedding_num_tokens", + "text-embedding", + {"texts": ["hello"]}, + [{"num_tokens": [5]}], + ), + ( + "invoke_rerank", + "rerank", + { + "query": "search query", + "docs": ["first document", "second document"], + "score_threshold": 0.5, + "top_n": 1, + }, + [ + { + "model": "echo", + "docs": [{"index": 0, "text": "first document", "score": 0.9}], + } + ], + ), + ( + "invoke_multimodal_rerank", + "rerank", + { + "query": {"content": "search query", "content_type": "text"}, + "docs": [{"content": "first document", "content_type": "text"}], + "score_threshold": 0.5, + "top_n": 1, + }, + [ + { + "model": "echo", + "docs": [{"index": 0, "text": "first document", "score": 0.9}], + } + ], + ), + ( + "invoke_tts", + "tts", + {"content_text": "hello", "voice": "echo", "tenant_id": "tenant"}, + [{"result": b"hello".hex(), "mime_type": "audio/wav"}], + ), + ( + "get_tts_model_voices", + "tts", + {"language": "en_US"}, + [{"voices": [{"name": "Echo", "value": "echo"}]}], + ), + ( + "invoke_speech2text", + "speech2text", + {"file": b"hello".hex()}, + [{"result": "hello"}], + ), + ("invoke_moderation", "moderation", {"text": "blocked"}, [{"result": True}]), + ( + "get_ai_model_schemas", + "llm", + {}, + [{"model_schema": {"model": "echo", "model_type": "llm"}}], + ), +] + +MODEL_METHODS = { + "validate_model_credentials": "validate_credentials", + "invoke_llm": "_invoke", + "start_polling": "_start_polling", + "check_polling": "_check_polling", + "get_llm_num_tokens": "get_num_tokens", + "invoke_text_embedding": "_invoke", + "invoke_multimodal_embedding": "_invoke_multimodal", + "get_text_embedding_num_tokens": "get_num_tokens", + "invoke_rerank": "_invoke", + "invoke_multimodal_rerank": "_invoke_multimodal", + "invoke_tts": "_invoke", + "get_tts_model_voices": "get_tts_model_voices", + "invoke_speech2text": "_invoke", + "invoke_moderation": "_invoke", + "get_ai_model_schemas": "get_model_schema", +} +RUNTIME_METHODS = { + "invoke_tool": ("get_tool_cls", "_invoke", "tool_parameters"), + "get_tool_runtime_parameters": ("get_tool_cls", "_get_runtime_parameters", None), + "invoke_agent_strategy": ( + "get_agent_strategy_cls", + "_invoke", + "agent_strategy_params", + ), + "invoke_website_datasource_get_crawl": ( + "get_website_crawl_datasource_cls", + "_get_website_crawl", + "datasource_parameters", + ), + "invoke_online_document_datasource_get_pages": ( + "get_online_document_datasource_cls", + "_get_pages", + "datasource_parameters", + ), + "invoke_online_document_datasource_get_page_content": ( + "get_online_document_datasource_cls", + "_get_content", + "page", + ), + "invoke_online_drive_browse_files": ( + "get_online_drive_datasource_cls", + "_browse_files", + "request", + ), + "invoke_online_drive_download_file": ( + "get_online_drive_datasource_cls", + "_download_file", + "request", + ), +} +TRIGGER_METHODS = { + "get_authorization_url": "_oauth_get_authorization_url", + "get_credentials": "_oauth_get_credentials", + "refresh_credentials": "_oauth_refresh_credentials", + "invoke_trigger_event": "_on_event", + "validate_trigger_credentials": "_validate_api_key", + "dispatch_trigger_event": "_dispatch_event", + "subscribe_trigger": "_create_subscription", + "unsubscribe_trigger": "_delete_subscription", + "refresh_trigger": "_refresh_subscription", + "fetch_parameter_options": "_fetch_parameter_options", +} +CASES += [ + ( + "model", + action, + { + "provider": "contract_models", + "model_type": model_type, + "model": "echo", + **payload, + }, + expected, + ) + for action, model_type, payload, expected in MODEL_CASES +] + + +def assert_contains(actual: object, expected: object) -> None: + if isinstance(expected, dict): + assert isinstance(actual, dict), actual + for key, value in expected.items(): + assert key in actual, actual + assert_contains(actual[key], value) + elif isinstance(expected, list): + assert isinstance(actual, list), actual + assert len(actual) == len(expected), actual + for item, expected_item in zip(actual, expected, strict=True): + assert_contains(item, expected_item) + else: + assert actual == expected + + +def test_every_public_action_has_one_route(contract_plugin: Plugin) -> None: + assert set(ACTIONS) == set(PluginInvokeType) + assert {(kind, action) for kind, action, _, _ in CASES} == { + (kind.value, action.value) + for kind, actions in ACTIONS.items() + for action in actions + } + for invoke_type, actions in ACTIONS.items(): + for action in actions: + request = {"type": invoke_type.value, "action": action.value} + assert ( + sum(route.filter(request) for route in contract_plugin.routes) == 1 + ), request + + +@pytest.mark.parametrize( + ("kind", "action", "payload", "expected"), + CASES, + ids=[f"{kind}.{action}" for kind, action, _, _ in CASES], +) +def test_public_actions_execute_registered_plugins( + contract_plugin: Plugin, + mocker: MockerFixture, + kind: str, + action: str, + payload: dict, + expected: list[dict], +) -> None: + request = deepcopy({ + "type": kind, + "action": action, + "user_id": "user-1", + "credentials": CREDENTIALS, + **payload, + }) + original = deepcopy(request) + if action in MODEL_METHODS: + model = contract_plugin.registration.get_model_instance( + payload["provider"], ModelType(payload["model_type"]) + ) + method = MODEL_METHODS[action] + parameters = signature(getattr(type(model), method)) + call = mocker.spy(type(model), method) + forwarded = deepcopy(request) + for key in ("type", "action", "provider", "model_type", "user_id"): + forwarded.pop(key) + if action.startswith("invoke_") or action in {"start_polling", "check_polling"}: + forwarded["user"] = "user-1" + elif action in RUNTIME_METHODS: + getter, method, parameter = RUNTIME_METHODS[action] + provider_cls = getattr(contract_plugin.registration, getter)( + request.get("agent_strategy_provider", request.get("provider")), + request.get( + "tool", request.get("datasource", request.get("agent_strategy")) + ), + ) + parameters = signature(getattr(provider_cls, method)) + call = mocker.spy(provider_cls, method) + forwarded = ( + { + "parameters" if kind == "agent_strategy" else parameter: deepcopy( + request[parameter] + ) + } + if parameter is not None + else {} + ) + elif action in TRIGGER_METHODS: + factory = contract_plugin.registration.trigger_factory + if action == "invoke_trigger_event": + _, provider_cls = factory.iter_events(request["provider"])[request["event"]] + elif action == "dispatch_trigger_event": + provider_cls = factory.get_provider_cls(request["provider"]) + else: + provider_cls = factory.get_subscription_constructor_cls(request["provider"]) + method = TRIGGER_METHODS[action] + parameters = signature(getattr(provider_cls, method)) + call = mocker.spy(provider_cls, method) + forwarded = { + name: deepcopy(request[name]) + for name in parameters.parameters + if name not in {"self", "request"} + } + if action == "get_tool_runtime_parameters": + with pytest.warns(DeprecationWarning, match="replaced by dynamic select"): + messages = invoke(contract_plugin, request) + else: + messages = invoke(contract_plugin, request) + assert len(messages) == len(expected), messages + for message, fields in zip(messages, expected, strict=True): + assert message["type"] == "stream", message + assert message["data"] + assert_contains(message["data"], fields) + + if ( + action in MODEL_METHODS + or action in RUNTIME_METHODS + or action in TRIGGER_METHODS + ): + call.assert_called_once() + received = parameters.bind(*call.call_args.args, **call.call_args.kwargs) + provider = received.arguments.pop("self") + if action in RUNTIME_METHODS: + assert provider.runtime.user_id == original["user_id"] + if kind != "agent_strategy": + assert provider.runtime.credentials == original["credentials"] + assert provider.runtime.session_id == "session-1" + if kind == "tool": + assert provider.runtime.credential_type.value == original.get( + "credential_type", "api-key" + ) + assert provider.session.session_id == "session-1" + elif action in TRIGGER_METHODS: + if kind != "oauth": + assert provider.runtime.credentials == original["credentials"] + assert provider.runtime.credential_type.value == original.get( + "credential_type", "api-key" + ) + assert provider.runtime.session.session_id == "session-1" + if action == "invoke_trigger_event": + assert ( + provider.runtime.subscription.model_dump() + == original["subscription"] + ) + if "request" in received.arguments: + http_request = received.arguments.pop("request") + assert ( + http_request.get_data() + == bytes.fromhex(original["raw_http_request"]).split( + b"\r\n\r\n", 1 + )[1] + ) + if action == "invoke_speech2text": + assert received.arguments.pop("file").closed + forwarded.pop("file") + actual = TypeAdapter(dict[str, Any]).dump_python( + received.arguments, mode="json", exclude_unset=True + ) + assert actual == forwarded + + +def test_tool_runtime_parameters_reach_the_plugin(contract_plugin: Plugin) -> None: + # Dify's ToolClient.get_runtime_parameters consumes this daemon action. + request = { + "type": "tool", + "action": "get_tool_runtime_parameters", + "user_id": "user-1", + "provider": "contract", + "tool": "echo", + "credentials": {"api_key": "valid"}, + } + with pytest.warns(DeprecationWarning, match="replaced by dynamic select"): + messages = invoke(contract_plugin, request) + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + parameters = messages[0]["data"]["parameters"] + assert len(parameters) == 2 + assert parameters[0]["name"] == "message" + assert parameters[0]["type"] == "string" + assert parameters[0]["required"] is True + + +def request_for(action: str, **overrides: object) -> dict: + kind, _, payload, _ = next(case for case in CASES if case[1] == action) + return deepcopy({ + "type": kind, + "action": action, + "user_id": "user-1", + "credentials": CREDENTIALS, + **payload, + **overrides, + }) + + +@pytest.mark.parametrize("action", [case[1] for case in CASES]) +def test_invalid_action_payload_reports_one_error_and_ends( + contract_plugin: Plugin, action: str +) -> None: + request = request_for(action) + required_field = "raw_http_request" if action == "invoke_endpoint" else "user_id" + del request[required_field] + + messages = invoke(contract_plugin, request) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert messages[0]["data"]["error_type"] == "ValidationError" + assert required_field in messages[0]["data"]["message"] + + +@pytest.mark.parametrize("action", [case[1] for case in CASES if case[0] != "endpoint"]) +def test_missing_provider_reports_one_error_and_ends( + contract_plugin: Plugin, action: str +) -> None: + field = ( + "agent_strategy_provider" if action == "invoke_agent_strategy" else "provider" + ) + messages = invoke( + contract_plugin, request_for(action, **{field: "missing-provider"}) + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert messages[0]["data"]["error_type"] in {"ValueError", "TypeError"} + assert "missing-provider" in messages[0]["data"]["message"] + + +@pytest.mark.parametrize( + "action", + [ + "validate_tool_credentials", + "validate_provider_credentials", + "validate_model_credentials", + "validate_datasource_credentials", + "validate_trigger_credentials", + ], +) +def test_invalid_credentials_preserve_plugin_failure( + contract_plugin: Plugin, action: str +) -> None: + messages = invoke( + contract_plugin, request_for(action, credentials={"api_key": "invalid"}) + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert messages[0]["data"]["error_type"] == "ValueError" + assert messages[0]["data"]["message"] == "invalid credentials" + + +@pytest.mark.parametrize( + ("action", "field"), + [ + ("invoke_tool", "tool"), + ("get_tool_runtime_parameters", "tool"), + ("invoke_agent_strategy", "agent_strategy"), + ("invoke_website_datasource_get_crawl", "datasource"), + ("invoke_online_document_datasource_get_pages", "datasource"), + ("invoke_online_drive_browse_files", "datasource"), + ("invoke_trigger_event", "event"), + ], +) +def test_missing_implementation_reports_the_requested_name( + contract_plugin: Plugin, action: str, field: str +) -> None: + messages = invoke(contract_plugin, request_for(action, **{field: "missing-action"})) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert "missing-action" in messages[0]["data"]["message"] + + +@pytest.mark.parametrize("size", [0, 1, 8191, 8192, 8193, 16384, 4 * 1024 * 1024]) +def test_blob_chunk_boundaries_round_trip( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch, size: int +) -> None: + content = bytes(range(256)) * (size // 256) + bytes(range(size % 256)) + metadata = {"mime_type": "application/octet-stream", "filename": "result.bin"} + + def invoke_blob(tool: Tool, parameters: dict) -> Generator[ToolInvokeMessage]: + del parameters + yield tool.create_blob_message(content, meta=metadata) + + plugin_cls = contract_plugin.registration.get_tool_cls("contract", "echo") + monkeypatch.setattr(plugin_cls, "_invoke", invoke_blob) + messages = invoke(contract_plugin, request_for("invoke_tool")) + + assert len(messages) == (size + 8191) // 8192 + 1 + chunks = [message["data"]["message"] for message in messages] + assert all(message["type"] == "stream" for message in messages) + assert all(message["data"]["type"] == "blob_chunk" for message in messages) + assert all(message["data"]["meta"] == metadata for message in messages) + assert len({chunk["id"] for chunk in chunks}) == 1 + assert len(chunks[0]["id"]) == 32 + assert [chunk["sequence"] for chunk in chunks] == list(range(len(chunks))) + assert all(chunk["total_length"] == size for chunk in chunks) + assert [chunk["end"] for chunk in chunks] == [False] * (len(chunks) - 1) + [True] + assert not chunks[-1]["blob"] + assert b"".join(base64.b64decode(chunk["blob"]) for chunk in chunks) == content + # The daemon rejects JSON lines over 5 MiB before Dify sees the message. + assert all(len(chunk["blob"]) <= 4 * ((8192 + 2) // 3) for chunk in chunks) + + +def test_tool_blob_subclasses_keep_the_base_chunk_contract( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch +) -> None: + class TaggedMessage(ToolInvokeMessage): + source: str + + def provider(*_args: object, **_kwargs: object) -> Generator[ToolInvokeMessage]: + yield TaggedMessage( + type=ToolInvokeMessage.MessageType.BLOB, + message=ToolInvokeMessage.BlobMessage(blob=b"custom"), + source="custom-tool", + ) + + plugin_cls = contract_plugin.registration.get_tool_cls("contract", "echo") + monkeypatch.setattr(plugin_cls, "_invoke", provider) + + messages = invoke(contract_plugin, request_for("invoke_tool")) + + assert len(messages) == 2 + assert all(message["type"] == "stream" for message in messages) + assert all( + set(message["data"]) == {"type", "message", "meta"} for message in messages + ) + assert all(message["data"]["type"] == "blob_chunk" for message in messages) + assert base64.b64decode(messages[0]["data"]["message"]["blob"]) == b"custom" + assert not messages[1]["data"]["message"]["blob"] + assert messages[1]["data"]["message"]["end"] is True + + +def test_repeated_tool_blob_keeps_existing_id_replacement( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch +) -> None: + original = ToolInvokeMessage( + type=ToolInvokeMessage.MessageType.BLOB, + message=ToolInvokeMessage.BlobMessage(blob=b"payload"), + ) + + def provider(*_args: object, **_kwargs: object) -> Generator[ToolInvokeMessage]: + yield original + yield original + + plugin_cls = contract_plugin.registration.get_tool_cls("contract", "echo") + monkeypatch.setattr(plugin_cls, "_invoke", provider) + + messages = invoke(contract_plugin, request_for("invoke_tool")) + + assert len(messages) == 4 + chunks = [message["data"]["message"] for message in messages] + assert base64.b64decode(chunks[0]["blob"]) == b"payload" + assert base64.b64decode(chunks[2]["blob"]) == chunks[0]["id"].encode() + assert [chunk["total_length"] for chunk in chunks] == [7, 7, 32, 32] + assert [chunk["end"] for chunk in chunks] == [False, True, False, True] + assert original.message.blob == chunks[2]["id"].encode() + + +@pytest.mark.parametrize("kind", ["drive", "document", "agent"]) +@pytest.mark.parametrize( + ("size", "metadata"), + [ + (0, None), + (4, {"file_name": "result.txt"}), + (4 * 1024 * 1024, {"mime_type": "application/octet-stream"}), + (4, {"file_name": "image.png", "mime_type": "image/png"}), + ], + ids=["empty", "named", "large", "image"], +) +def test_datasource_and_agent_blobs_preserve_wire_format_without_uploads( + contract_plugin: Plugin, + mocker: MockerFixture, + kind: str, + size: int, + metadata: dict | None, +) -> None: + content = bytes(range(256)) * (size // 256) + bytes(range(size % 256)) + upload = mocker.patch.object(File, "upload") + message_type = AgentInvokeMessage if kind == "agent" else DatasourceMessage + original = message_type( + type=message_type.MessageType.BLOB, + message=message_type.BlobMessage(blob=content), + meta=metadata, + ) + expected = original.model_dump(mode="json") + finalized = [] + + def provider(*_args: object, **_kwargs: object) -> Generator: + try: + yield original + finally: + finalized.append(True) + + if kind == "drive": + plugin_cls = contract_plugin.registration.get_online_drive_datasource_cls( + "drive", "echo" + ) + method, action = "_download_file", "invoke_online_drive_download_file" + elif kind == "document": + plugin_cls = contract_plugin.registration.get_online_document_datasource_cls( + "documents", "echo" + ) + method, action = ( + "_get_content", + "invoke_online_document_datasource_get_page_content", + ) + else: + plugin_cls = contract_plugin.registration.get_agent_strategy_cls( + "contract_agents", "echo" + ) + method, action = "_invoke", "invoke_agent_strategy" + mocker.patch.object(plugin_cls, method, provider) + + messages = invoke(contract_plugin, request_for(action)) + + assert messages == [{"type": "stream", "data": expected}] + assert messages[0]["data"]["type"] == "blob" + assert base64.b64decode(messages[0]["data"]["message"]["blob"]) == content + upload.assert_not_called() + assert finalized == [True] + + +def test_partial_tool_stream_preserves_data_then_reports_failure( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch +) -> None: + finalized = [] + + def fail(tool: Tool, parameters: dict) -> Generator[ToolInvokeMessage]: + del parameters + try: + yield tool.create_text_message("first") + message = "stream interrupted" + raise ValueError(message) + finally: + finalized.append(True) + + tool_cls = contract_plugin.registration.get_tool_cls("contract", "echo") + monkeypatch.setattr(tool_cls, "_invoke", fail) + messages = invoke(contract_plugin, request_for("invoke_tool")) + + assert [message["type"] for message in messages] == ["stream", "error"] + assert messages[0]["data"]["message"] == {"text": "first"} + assert messages[1]["data"]["message"] == "stream interrupted" + assert finalized == [True] + + +def test_endpoint_response_includes_route_values_and_settings( + contract_plugin: Plugin, +) -> None: + messages = invoke(contract_plugin, request_for("invoke_endpoint")) + + assert messages[0]["data"]["headers"]["X-Item"] == "item" + assert messages[0]["data"]["headers"]["X-Setting"] == "setting" + + +@pytest.mark.parametrize("path", ["/missing", "/echo/item"]) +def test_endpoint_missing_path_or_method_returns_http_404( + contract_plugin: Plugin, path: str +) -> None: + messages = invoke( + contract_plugin, + request_for( + "invoke_endpoint", raw_http_request=raw_request(path, method="GET") + ), + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert messages[0]["data"]["status"] == 404 + assert b"Failed to dispatch endpoint request" in bytes.fromhex( + messages[0]["data"]["result"] + ) + + +def test_endpoint_exception_returns_http_500(contract_plugin: Plugin) -> None: + messages = invoke(contract_plugin, request_for("invoke_endpoint", settings={})) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert messages[0]["data"]["status"] == 500 + assert b"Internal Server Error" in bytes.fromhex(messages[0]["data"]["result"]) + + +def test_endpoint_stream_emits_headers_before_utf8_and_binary_chunks( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch +) -> None: + endpoint_cls = next(contract_plugin.registration.endpoints.iter_rules()).endpoint + + def response(*args: object) -> Response: + del args + + def chunks() -> Generator[str | bytes]: + yield "你好" + yield b"\x00\xff" + + return Response(chunks(), status=206, headers={"X-Test": "stream"}) + + monkeypatch.setattr(endpoint_cls, "_invoke", response) + messages = invoke(contract_plugin, request_for("invoke_endpoint")) + + assert len(messages) == 3 + assert all(message["type"] == "stream" for message in messages) + assert messages[0]["data"]["status"] == 206 + assert messages[0]["data"]["headers"]["X-Test"] == "stream" + assert "result" not in messages[0]["data"] + assert bytes.fromhex(messages[1]["data"]["result"]).decode() == "你好" + assert bytes.fromhex(messages[2]["data"]["result"]) == b"\x00\xff" + + +@pytest.mark.parametrize("credential_type", ["api-key", "oauth2", "unauthorized"]) +def test_dynamic_tool_options_preserve_runtime_credential_type_and_identity( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch, credential_type: str +) -> None: + runtimes = [] + + def options(tool: Tool, parameter: str) -> list[ParameterOption]: + runtimes.append(tool.runtime) + return [ + ParameterOption( + value=tool.runtime.credential_type.value, + label=I18nObject(en_us=parameter), + ) + ] + + tool_cls = contract_plugin.registration.get_tool_cls("contract", "echo") + monkeypatch.setattr(tool_cls, "_fetch_parameter_options", options) + messages = invoke( + contract_plugin, + request_for( + "fetch_parameter_options", + provider="contract", + provider_action="echo", + credential_type=credential_type, + ), + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert_contains( + messages[0]["data"], + {"options": [{"value": "api-key", "label": {"en_US": "event"}}]}, + ) + assert len(runtimes) == 1 + assert runtimes[0].credentials == CREDENTIALS + assert runtimes[0].user_id == "user-1" + assert runtimes[0].session_id == "session-1" + + +def test_dynamic_event_options_dispatch_to_the_event( + contract_plugin: Plugin, mocker: MockerFixture +) -> None: + _, event_cls = contract_plugin.registration.trigger_factory.iter_events("events")[ + "echo" + ] + call = mocker.spy(event_cls, "_fetch_parameter_options") + messages = invoke( + contract_plugin, + request_for( + "fetch_parameter_options", provider_action="echo", credential_type="oauth2" + ), + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert_contains( + messages[0]["data"], + {"options": [{"value": "echo", "label": {"en_US": "event"}}]}, + ) + call.assert_called_once() + runtime = call.call_args.args[0].runtime + assert runtime.credentials == CREDENTIALS + assert runtime.credential_type.value == "oauth2" + assert runtime.session.session_id == "session-1" + + +def test_unknown_dynamic_action_returns_an_error(contract_plugin: Plugin) -> None: + messages = invoke( + contract_plugin, + request_for("fetch_parameter_options", provider_action="missing"), + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert ( + messages[0]["data"]["message"] + == "Cannot find the target to fetch parameter options" + ) + + +@pytest.mark.parametrize("provider", ["contract", "website"]) +def test_oauth_selects_tool_and_datasource_providers( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch, provider: str +) -> None: + registration = contract_plugin.registration + schema = registration.trigger_factory.get_configuration( + "events" + ).subscription_constructor.oauth_schema + if provider == "contract": + configuration, provider_cls, _ = registration.tools_mapping[provider] + else: + mapping = registration.datasource_mapping[provider] + configuration, provider_cls = mapping.configuration, mapping.provider_cls + configuration.oauth_schema = schema + + def authorization_url( + self: object, redirect_uri: str, system_credentials: dict + ) -> str: + assert isinstance(self, provider_cls) + assert system_credentials == {"client_id": "client"} + return redirect_uri + "?provider=" + provider + + monkeypatch.setattr(provider_cls, "_oauth_get_authorization_url", authorization_url) + messages = invoke( + contract_plugin, request_for("get_authorization_url", provider=provider) + ) + + assert messages == [ + { + "type": "stream", + "data": { + "authorization_url": "https://plugin.invalid/callback?provider=" + + provider + }, + } + ] + + +def test_trigger_ignore_is_a_cancelled_result(contract_plugin: Plugin) -> None: + messages = invoke( + contract_plugin, + request_for("invoke_trigger_event", parameters={"ignore": True}), + ) + + assert messages == [ + {"type": "stream", "data": {"variables": {}, "cancelled": True}} + ] + + +@pytest.mark.parametrize("framing", ["content-length", "decoded-chunked"]) +@pytest.mark.parametrize( + ("content_type", "body", "expected"), + [ + ( + "application/x-www-form-urlencoded", + b"payload=%7B%22message%22%3A%22caf%C3%A9%22%7D&tag=one&tag=two&empty=&plus=a%2Bb", + { + "payload": ['{"message":"café"}'], + "tag": ["one", "two"], + "empty": [""], + "plus": ["a+b"], + }, + ), + ( + "application/json", + b'{\n "message": "caf\\u00e9", "id": 9007199254740993\n}', + {"message": "café", "id": 9007199254740993}, + ), + ], + ids=["form", "json"], +) +def test_webhook_dispatch_preserves_signed_body_and_request_fields( + contract_plugin: Plugin, + monkeypatch: pytest.MonkeyPatch, + framing: str, + content_type: str, + body: bytes, + expected: dict, +) -> None: + signature = hmac.digest(b"webhook-secret", body, "sha256").hex() + framing_header = ( + f"Content-Length: {len(body)}" + if framing == "content-length" + else "Transfer-Encoding: chunked" + ) + # Dify serializes get_data() after HTTP transfer decoding, retaining headers. + raw = ( + "POST /events/你好?state=a%2Bb&tag=one&tag=two HTTP/1.1\r\n" + "Host: plugin.invalid\r\n" + f"Content-Type: {content_type}\r\n{framing_header}\r\n" + f"X-Webhook-Signature: {signature}\r\n\r\n" + ).encode() + body + provider_cls = contract_plugin.registration.trigger_factory.get_provider_cls( + "events" + ) + + def dispatch( + _self: object, subscription: Subscription, request: Request + ) -> EventDispatch: + assert subscription.endpoint == SUBSCRIPTION["endpoint"] + assert request.method == "POST" + assert request.path == "/events/你好" + assert request.query_string == b"state=a%2Bb&tag=one&tag=two" + assert request.args.to_dict(flat=False) == { + "state": ["a+b"], + "tag": ["one", "two"], + } + raw_body = request.get_data(cache=True) + assert raw_body == body + assert hmac.compare_digest( + request.headers["X-Webhook-Signature"], + hmac.digest(b"webhook-secret", raw_body, "sha256").hex(), + ) + payload = ( + request.get_json() + if content_type == "application/json" + else request.form.to_dict(flat=False) + ) + assert payload == expected + return EventDispatch( + events=["echo"], response=Response("accepted", status=202), payload=payload + ) + + monkeypatch.setattr(provider_cls, "_dispatch_event", dispatch) + messages = invoke( + contract_plugin, + request_for("dispatch_trigger_event", raw_http_request=raw.hex()), + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert messages[0]["data"]["events"] == ["echo"] + assert messages[0]["data"]["payload"] == expected + response = bytes.fromhex(messages[0]["data"]["response"]) + head, response_body = response.split(b"\r\n\r\n", 1) + assert head.startswith(b"HTTP/1.1 202 ") + assert response_body == b"accepted" + + +@pytest.mark.parametrize("file_body", [False, True]) +def test_trigger_dispatch_closes_response_resources( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch, file_body: bool +) -> None: + resource = BytesIO(b"accepted") + response = Response( + FileWrapper(resource) if file_body else resource.getvalue(), status=202 + ) + if not file_body: + response.call_on_close(resource.close) + provider_cls = contract_plugin.registration.trigger_factory.get_provider_cls( + "events" + ) + monkeypatch.setattr( + provider_cls, + "_dispatch_event", + lambda *_args, **_kwargs: EventDispatch(events=["echo"], response=response), + ) + + try: + messages = invoke(contract_plugin, request_for("dispatch_trigger_event")) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert messages[0]["data"]["events"] == ["echo"] + raw = bytes.fromhex(messages[0]["data"]["response"]) + assert raw.startswith(b"HTTP/1.1 202 ") + assert raw.split(b"\r\n\r\n", 1)[1] == b"accepted" + assert resource.closed + finally: + response.close() + resource.close() + + +@pytest.mark.parametrize( + "action", + [ + "invoke_endpoint", + "get_credentials", + "invoke_trigger_event", + "dispatch_trigger_event", + ], +) +def test_invalid_http_encoding_reports_an_error( + contract_plugin: Plugin, action: str +) -> None: + messages = invoke(contract_plugin, request_for(action, raw_http_request="not-hex!")) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert messages[0]["data"]["error_type"] == "Error" + assert "Non-hexadecimal" in messages[0]["data"]["message"] + + +def test_llm_non_streaming_result_preserves_message(contract_plugin: Plugin) -> None: + messages = invoke(contract_plugin, request_for("invoke_llm", stream=False)) + + assert len(messages) == 1 + assert messages[0]["type"] == "stream" + assert_contains( + messages[0]["data"], + { + "model": "echo", + "delta": {"message": {"role": "assistant", "content": "hello"}}, + }, + ) + + +@pytest.mark.parametrize("provider", ["contract", "website"]) +def test_registered_provider_without_oauth_reports_unsupported( + contract_plugin: Plugin, provider: str +) -> None: + registration = contract_plugin.registration + if provider == "contract": + registration.tools_mapping[provider][0].oauth_schema = None + else: + registration.datasource_mapping[provider].configuration.oauth_schema = None + messages = invoke( + contract_plugin, request_for("get_authorization_url", provider=provider) + ) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert ( + messages[0]["data"]["message"] + == f"Provider `{provider}` does not support OAuth" + ) + + +def test_tool_without_runtime_parameters_reports_unsupported( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch +) -> None: + tool_cls = contract_plugin.registration.get_tool_cls("contract", "echo") + monkeypatch.delattr(tool_cls, "_get_runtime_parameters") + + messages = invoke(contract_plugin, request_for("get_tool_runtime_parameters")) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert ( + messages[0]["data"]["message"] + == "Tool `echo` does not implement runtime parameters" + ) + + +@pytest.mark.parametrize("action", ["start_polling", "check_polling"]) +def test_model_without_polling_reports_unsupported_before_invocation( + contract_plugin: Plugin, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + model = contract_plugin.registration.get_model_instance( + "contract_models", ModelType.LLM + ) + monkeypatch.delattr(type(model), "_start_polling") + + messages = invoke(contract_plugin, request_for(action)) + + assert len(messages) == 1 + assert messages[0]["type"] == "error" + assert "does not support polling" in messages[0]["data"]["message"] diff --git a/tests/core/test_position_helper.py b/tests/core/test_position_helper.py new file mode 100644 index 00000000..3fa91642 --- /dev/null +++ b/tests/core/test_position_helper.py @@ -0,0 +1,67 @@ +from operator import itemgetter +from pathlib import Path + +import pytest + +from dify_plugin.core.utils.position_helper import ( + get_position_map, + sort_by_position_map, + sort_to_dict_by_position_map, +) + + +@pytest.mark.parametrize("content", [None, "", "[]", "[unterminated"]) +def test_absent_empty_or_invalid_positions_mean_no_order( + tmp_path: Path, content: str | None +) -> None: + if content is not None: + (tmp_path / "_position.yaml").write_text(content, encoding="utf-8") + + assert get_position_map(str(tmp_path)) == {} + + +def test_position_file_preserves_legacy_indices_after_trimming(tmp_path: Path) -> None: + (tmp_path / "custom.yaml").write_text( + '- " beta "\n- null\n- 42\n- ""\n- " "\n- alpha\n', encoding="utf-8" + ) + + assert get_position_map(str(tmp_path), file_name="custom.yaml") == { + "beta": 0, + "": 1, + "alpha": 2, + } + + +def test_sort_preserves_unlisted_order_without_mutating_input() -> None: + items = [{"name": name} for name in ("unknown-b", "alpha", "beta", "unknown-a")] + original = items.copy() + + ordered = sort_by_position_map({"beta": 0, "alpha": 1}, items, itemgetter("name")) + + assert [item["name"] for item in ordered] == [ + "beta", + "alpha", + "unknown-b", + "unknown-a", + ] + assert items == original + assert ordered[0] is items[2] + + +@pytest.mark.parametrize( + ("positions", "items"), [({}, ["beta", "alpha"]), ({"alpha": 0}, [])] +) +def test_sort_without_positions_or_items_preserves_values( + positions: dict[str, int], items: list[str] +) -> None: + assert sort_by_position_map(positions, items, str) == items + + +def test_ordered_mapping_keeps_names_and_original_objects() -> None: + items = [{"name": "alpha", "value": 1}, {"name": "beta", "value": 2}] + + ordered = sort_to_dict_by_position_map({"beta": 0}, items, itemgetter("name")) + + assert list(ordered) == ["beta", "alpha"] + assert ordered["beta"] is items[1] + assert ordered["alpha"] is items[0] diff --git a/tests/core/test_runtime.py b/tests/core/test_runtime.py new file mode 100644 index 00000000..3395c3b4 --- /dev/null +++ b/tests/core/test_runtime.py @@ -0,0 +1,612 @@ +import json +from collections.abc import Callable, Generator, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import closing + +import gevent +import httpx +import pytest +from gevent.event import Event +from pydantic import BaseModel, ValidationError + +from dify_plugin.config.config import InstallMethod +from dify_plugin.core.entities.invocation import InvokeType +from dify_plugin.core.entities.plugin.io import PluginInStream, PluginInStreamEvent +from dify_plugin.core.runtime import BackwardsInvocation, Session +from tests.http_server import Response, TestHTTPServer +from tests.servers.helpers import MemoryRequestReader, RecordingResponseWriter + + +class Result(BaseModel): + answer: str + + +class DaemonWriter(RecordingResponseWriter): + def __init__(self, on_request: Callable[[dict], None]) -> None: + super().__init__() + self.on_request = on_request + + def write(self, data: str) -> None: + super().write(data) + if data.strip(): + self.on_request(json.loads(data)) + + +@pytest.fixture +def session() -> Generator[Session]: + reader = MemoryRequestReader() + with ThreadPoolExecutor(max_workers=1) as executor: + session = Session( + session_id="session-a", + executor=executor, + reader=reader, + writer=RecordingResponseWriter(), + install_method=InstallMethod.Local, + dify_plugin_daemon_url="http://daemon.example", + max_invocation_timeout=2, + ) + try: + yield session + finally: + reader.close() + + +@pytest.fixture +def invocation( + session: Session, monkeypatch: pytest.MonkeyPatch +) -> BackwardsInvocation[Result]: + invocation = BackwardsInvocation[Result](session) + monkeypatch.setattr( + invocation, "_generate_backwards_request_id", lambda: "request-a" + ) + return invocation + + +def _response( + session: Session, + data: dict | None = None, + *, + request_id: str = "request-a", + event: str = "response", + message: str = "", +) -> PluginInStream: + # daemon 635d262e: session_manager/session.go and backwards_invocation/entities.go. + return PluginInStream( + session_id=session.session_id, + event=PluginInStreamEvent.BackwardInvocationResponse, + data={ + "backwards_request_id": request_id, + "event": event, + "message": message, + "data": data, + }, + reader=session.reader, + writer=session.writer, + ) + + +def _wire(chunk: PluginInStream) -> bytes: + return ( + json.dumps( + { + "session_id": chunk.session_id, + "event": chunk.event.value, + "data": chunk.data, + }, + ensure_ascii=False, + ).encode() + + b"\n\n" + ) + + +def _chunks(*chunks: PluginInStream | None) -> Generator[PluginInStream | None]: + yield from chunks + + +class ResponseStream(httpx.SyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self.chunks = chunks + self.closed = False + + def __iter__(self) -> Iterator[bytes]: + yield from self.chunks + + def close(self) -> None: + self.closed = True + + +def _http_client( + session: Session, + monkeypatch: pytest.MonkeyPatch, + *chunks: bytes, +) -> tuple[httpx.Client, ResponseStream]: + session.install_method = InstallMethod.Serverless + stream = ResponseStream(chunks) + client = httpx.Client( + transport=httpx.MockTransport( + lambda _request: httpx.Response(200, stream=stream) + ) + ) + monkeypatch.setattr("dify_plugin.core.runtime.httpx.Client", lambda: client) + return client, stream + + +@pytest.mark.parametrize("install_method", [InstallMethod.Local, InstallMethod.Remote]) +def test_backwards_invoke_receives_immediate_response( + session: Session, + invocation: BackwardsInvocation[Result], + install_method: InstallMethod, +) -> None: + session.install_method = install_method + + def respond(request: dict) -> None: + assert request == { + "event": "session", + "session_id": "session-a", + "data": { + "type": "invoke", + "data": { + "type": "storage", + "backwards_request_id": "request-a", + "request": {"opt": "get", "key": "你好"}, + }, + }, + } + session.reader._process_line(_response(session, {"answer": "stored"})) + session.reader._process_line(_response(session, event="end")) + + session.writer = DaemonWriter(respond) + with gevent.Timeout(3): + result = list( + invocation._backwards_invoke( + InvokeType.Storage, Result, {"opt": "get", "key": "你好"} + ) + ) + assert result == [Result(answer="stored")] + assert session.reader.readers == [] + + +def test_backwards_requests_do_not_share_results(session: Session) -> None: + sent = Event() + + def record_request(_request: dict) -> None: + if len(writer.messages) == 2: + sent.set() + + writer = DaemonWriter(record_request) + session.writer = writer + other = Session( + session_id="session-b", + executor=session._executor, + reader=session.reader, + writer=session.writer, + install_method=InstallMethod.Remote, + max_invocation_timeout=2, + ) + + def invoke(active_session: Session) -> list[Result]: + return list( + BackwardsInvocation[Result](active_session)._backwards_invoke( + InvokeType.Storage, Result, {"key": active_session.session_id} + ) + ) + + workers = [gevent.spawn(invoke, active) for active in (session, other)] + try: + assert sent.wait(timeout=2), "backwards requests were not sent" + ids = { + frame["session_id"]: frame["data"]["data"]["backwards_request_id"] + for frame in writer.messages + } + assert len(set(ids.values())) == 2 + session.reader._process_line( + _response(session, {"answer": "unrelated"}, request_id="unknown") + ) + session.reader._process_line( + _response( + other, {"answer": "wrong session"}, request_id=ids[session.session_id] + ) + ) + session.reader._process_line( + PluginInStream( + session_id=session.session_id, + event=PluginInStreamEvent.Request, + data={"backwards_request_id": ids[session.session_id]}, + reader=session.reader, + writer=session.writer, + ) + ) + for active, answer in ( + (other, "b-first"), + (session, "a-first"), + (other, "b-last"), + ): + session.reader._process_line( + _response(active, {"answer": answer}, request_id=ids[active.session_id]) + ) + session.reader._process_line( + _response(other, event="end", request_id=ids[other.session_id]) + ) + session.reader._process_line( + _response(session, event="end", request_id=ids[session.session_id]) + ) + assert workers[0].get(timeout=2) == [Result(answer="a-first")] + assert workers[1].get(timeout=2) == [ + Result(answer="b-first"), + Result(answer="b-last"), + ] + assert session.reader.readers == [] + finally: + gevent.killall(workers, block=True, timeout=2) + + +def test_cancelling_waiting_invocation_removes_subscription(session: Session) -> None: + sent = Event() + session.writer = DaemonWriter(lambda _request: sent.set()) + worker = gevent.spawn( + lambda: list( + BackwardsInvocation(session)._backwards_invoke(InvokeType.Storage, dict, {}) + ) + ) + try: + assert sent.wait(timeout=2) + assert len(session.reader.readers) == 1 + worker.kill(block=True, timeout=2) + assert worker.dead + assert session.reader.readers == [] + finally: + worker.kill(block=True, timeout=2) + + +def test_send_failure_removes_subscription( + session: Session, invocation: BackwardsInvocation[Result] +) -> None: + class FailedWriter(RecordingResponseWriter): + def write(self, _data: str) -> None: + msg = "daemon disconnected" + raise BrokenPipeError(msg) + + session.writer = FailedWriter() + with pytest.raises(BrokenPipeError, match="daemon disconnected"): + list(invocation._backwards_invoke(InvokeType.Storage, Result, {})) + assert session.reader.readers == [] + + +@pytest.mark.parametrize("ending", ["end", "error", "close", "invalid-result"]) +def test_full_duplex_releases_subscription_on_every_exit( + session: Session, invocation: BackwardsInvocation[Result], ending: str +) -> None: + terminal = { + "end": _response(session, event="end"), + "error": _response(session, event="error", message="permission denied"), + "invalid-result": _response(session, {"answer": []}), + "close": _response(session, {"answer": "unconsumed"}), + }[ending] + + def respond(_request: dict) -> None: + session.reader._process_line(_response(session, {"answer": "first"})) + session.reader._process_line(terminal) + + session.writer = DaemonWriter(respond) + with closing( + invocation._backwards_invoke(InvokeType.Storage, Result, {}) + ) as stream: + assert next(stream) == Result(answer="first") + assert len(session.reader.readers) == 1 + if ending == "end": + assert list(stream) == [] + elif ending == "close": + stream.close() + else: + message = "permission denied" if ending == "error" else "Failed to parse" + with pytest.raises(Exception, match=message): + next(stream) + assert session.reader.readers == [] + + +def test_reader_disconnect_unblocks_waiting_backwards_invocation( + session: Session, invocation: BackwardsInvocation[Result] +) -> None: + sent = Event() + session.writer = DaemonWriter(lambda _request: sent.set()) + worker = gevent.spawn( + lambda: list(invocation._backwards_invoke(InvokeType.Storage, Result, {})) + ) + try: + assert sent.wait(timeout=2) + session.reader.close() + assert worker.get(timeout=2) == [] + assert session.reader.readers == [] + finally: + worker.kill(block=True, timeout=2) + + +def test_data_resets_consecutive_empty_response_timeout( + session: Session, invocation: BackwardsInvocation[Result] +) -> None: + result = invocation._line_converter_wrapper( + _chunks( + None, + _response(session, {"answer": "first"}), + None, + _response(session, {"answer": "second"}), + _response(session, event="end"), + _response(session, {"answer": "after end"}), + ), + Result, + ) + assert list(result) == [Result(answer="first"), Result(answer="second")] + + +def test_empty_response_count_preserves_timeout_exception( + session: Session, + invocation: BackwardsInvocation[Result], +) -> None: + session.max_invocation_timeout = 2 + with pytest.raises(Exception, match="without response after 2 seconds") as error: + list(invocation._line_converter_wrapper(_chunks(None, None), Result)) + assert type(error.value) is Exception + + +def test_invocation_timeout_resets_for_results_and_excludes_consumer_time( + session: Session, + invocation: BackwardsInvocation[Result], +) -> None: + session.max_invocation_timeout = 0.1 + + def respond() -> Generator[PluginInStream]: + for text in ("first", "second"): + gevent.sleep(0.06) + yield _response(session, {"answer": text}) + yield _response(session, event="end") + + stream = invocation._line_converter_wrapper(respond(), Result) + assert next(stream) == Result(answer="first") + gevent.sleep(0.15) + assert list(stream) == [Result(answer="second")] + + +@pytest.mark.parametrize("terminal_event", ["end", "response"]) +def test_null_response_terminates_without_constructing_result( + session: Session, + invocation: BackwardsInvocation[Result], + terminal_event: str, +) -> None: + assert ( + list( + invocation._line_converter_wrapper( + _chunks(_response(session, event=terminal_event)), Result + ) + ) + == [] + ) + + +def test_daemon_error_preserves_timeout_phase( + session: Session, invocation: BackwardsInvocation[Result] +) -> None: + message = "invoke llm failed: phase=read_idle timeout_ms=240000" + with pytest.raises(Exception, match=message): + list( + invocation._line_converter_wrapper( + _chunks(_response(session, event="error", message=message)), Result + ) + ) + + +def test_invalid_result_retains_validation_cause( + session: Session, invocation: BackwardsInvocation[Result] +) -> None: + with pytest.raises(Exception, match="Failed to parse response") as exc: + list( + invocation._line_converter_wrapper( + _chunks(_response(session, {"answer": []})), Result + ) + ) + assert isinstance(exc.value.__cause__, ValidationError) + + +def test_missing_session_cannot_start_backwards_invocation() -> None: + with pytest.raises(Exception, match="does not support backwards invoke"): + BackwardsInvocation()._backwards_invoke(InvokeType.Storage, dict, {}) + + +def test_serverless_invocation_requires_daemon_url(session: Session) -> None: + session.install_method = InstallMethod.Serverless + session.dify_plugin_daemon_url = None + with pytest.raises(Exception, match="does not support backwards invoke"): + list( + BackwardsInvocation(session)._backwards_invoke(InvokeType.Storage, dict, {}) + ) + + +@pytest.mark.parametrize( + "ending", ["end", "error", "invalid-json", "close", "consumer-error", "eof"] +) +def test_http_stream_is_closed_on_every_exit( + session: Session, + invocation: BackwardsInvocation[Result], + monkeypatch: pytest.MonkeyPatch, + ending: str, +) -> None: + end = { + "end": _wire(_response(session, event="end")), + "error": _wire(_response(session, event="error", message="permission denied")), + "invalid-json": b"not-json\n", + "eof": b"", + "close": b"", + "consumer-error": b"", + }[ending] + client, response_stream = _http_client( + session, monkeypatch, _wire(_response(session, {"answer": "first"})), end + ) + with closing( + invocation._backwards_invoke(InvokeType.Storage, Result, {}) + ) as stream: + assert next(stream) == Result(answer="first") + assert not response_stream.closed + if ending == "end": + assert list(stream) == [] + elif ending == "error": + with pytest.raises(Exception, match="permission denied"): + next(stream) + elif ending == "invalid-json": + with pytest.raises(ValidationError, match="Invalid JSON"): + next(stream) + elif ending == "eof": + assert list(stream) == [] + elif ending == "consumer-error": + with pytest.raises( + Exception, match="Failed to parse response: consumer stopped" + ) as error: + stream.throw(ValueError("consumer stopped")) + assert isinstance(error.value.__cause__, ValueError) + else: + stream.close() + assert response_stream.closed + assert client.is_closed + + +@pytest.mark.integration +def test_http_backwards_invocation_uses_daemon_wire_protocol( + session: Session, + invocation: BackwardsInvocation[Result], + http_server: Callable[[list[Response]], TestHTTPServer], +) -> None: + session.install_method = InstallMethod.Serverless + body = ( + b"\n\n" + + _wire(_response(session, {"answer": "你好"})) + + _wire(_response(session, event="end")) + ) + # A Unicode code point and the frame delimiter both cross HTTP chunk boundaries. + boundary = body.index("你".encode()) + 1 + server = http_server([ + Response( + headers={"Content-Type": "text/event-stream"}, + chunks=(body[:boundary], body[boundary:-1], body[-1:]), + ) + ]) + session.dify_plugin_daemon_url = server.url + "/prefix/" + assert list( + invocation._backwards_invoke( + InvokeType.Storage, Result, {"opt": "get", "key": "你好"} + ) + ) == [Result(answer="你好")] + assert len(server.requests) == 1 + request = server.requests[0] + assert request.method == "POST" + assert request.path == "/prefix/backwards-invocation/transaction" + assert request.headers["Dify-Plugin-Session-ID"] == session.session_id + assert request.body.endswith(b"\n\n") + assert json.loads(request.body) == { + "event": "session", + "session_id": "session-a", + "data": { + "type": "invoke", + "data": { + "type": "storage", + "backwards_request_id": "request-a", + "request": {"opt": "get", "key": "你好"}, + }, + }, + } + + +@pytest.mark.integration +@pytest.mark.parametrize( + ("status", "body", "error"), + [ + ( + 400, + ( + b'{"backwards_request_id":"request-a","event":"error",' + b'"message":"failed to get session info from cache",' + b'"data":{"error_type":"SessionNotFound"}}' + ), + KeyError, + ), + (401, b"", None), + (429, b"", None), + (500, b"failed to parse request", ValidationError), + ], +) +def test_http_status_preserves_response_parsing( + session: Session, + invocation: BackwardsInvocation[Result], + http_server: Callable[[list[Response]], TestHTTPServer], + status: int, + body: bytes, + error: type[Exception] | None, +) -> None: + session.install_method = InstallMethod.Serverless + session.dify_plugin_daemon_url = http_server([ + Response(status=status, body=body) + ]).url + stream = invocation._backwards_invoke(InvokeType.Storage, Result, {}) + if error is None: + assert list(stream) == [] + else: + with pytest.raises(error): + list(stream) + + +@pytest.mark.integration +def test_http_read_timeout_releases_response( + session: Session, + invocation: BackwardsInvocation[Result], + monkeypatch: pytest.MonkeyPatch, + http_server: Callable[[list[Response]], TestHTTPServer], +) -> None: + session.install_method = InstallMethod.Serverless + session.max_invocation_timeout = 1 + session.dify_plugin_daemon_url = http_server([ + Response(body=_wire(_response(session, event="end")), delay=2) + ]).url + responses: list[httpx.Response] = [] + client = httpx.Client(event_hooks={"response": [responses.append]}) + monkeypatch.setattr("dify_plugin.core.runtime.httpx.Client", lambda: client) + with gevent.Timeout(3), pytest.raises(httpx.ReadTimeout): + list(invocation._backwards_invoke(InvokeType.Storage, Result, {})) + assert len(responses) == 1 + assert responses[0].is_closed + assert client.is_closed + + +@pytest.mark.integration +def test_http_transport_failure_does_not_complete_partial_result( + session: Session, + invocation: BackwardsInvocation[Result], + http_server: Callable[[list[Response]], TestHTTPServer], +) -> None: + session.install_method = InstallMethod.Serverless + body = _wire(_response(session, {"answer": "partial"})) + server = http_server([ + Response(headers={"Content-Length": str(len(body) + 1)}, body=body) + ]) + session.dify_plugin_daemon_url = server.url + with closing( + invocation._backwards_invoke(InvokeType.Storage, Result, {}) + ) as stream: + assert next(stream) == Result(answer="partial") + with pytest.raises(httpx.RemoteProtocolError, match="complete message body"): + next(stream) + + +@pytest.mark.integration +def test_http_heartbeats_keep_the_stream_open( + session: Session, + invocation: BackwardsInvocation[Result], + http_server: Callable[[list[Response]], TestHTTPServer], +) -> None: + session.install_method = InstallMethod.Serverless + session.max_invocation_timeout = 0.1 + session.dify_plugin_daemon_url = http_server([ + Response( + chunks=(*([b"\n\n"] * 10), _wire(_response(session, event="end"))), + delay=0.03, + ) + ]).url + + with gevent.Timeout(2): + assert list(invocation._backwards_invoke(InvokeType.Storage, Result, {})) == [] diff --git a/tests/core/test_session_context.py b/tests/core/test_session_context.py index c7fb7aec..1f388ffc 100644 --- a/tests/core/test_session_context.py +++ b/tests/core/test_session_context.py @@ -1,20 +1,33 @@ from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor from unittest.mock import MagicMock +import gevent +import pytest +from gevent.event import Event + import dify_plugin -from dify_plugin.core.entities.plugin.request import ModelInvokeTextEmbeddingRequest +from dify_plugin.core.entities.plugin.request import ( + ModelInvokeLLMRequest, + ModelInvokeTextEmbeddingRequest, +) from dify_plugin.core.plugin_executor import PluginExecutor from dify_plugin.core.runtime import Session -from dify_plugin.core.session_context import get_current_session +from dify_plugin.core.session_context import get_current_session, use_current_session from dify_plugin.entities.model import EmbeddingInputType from dify_plugin.interfaces.model.large_language_model import LargeLanguageModel from dify_plugin.interfaces.model.text_embedding_model import TextEmbeddingModel +from tests.servers.helpers import MemoryRequestReader, RecordingResponseWriter def _session(app_id: str = "app-1") -> Session: - session = Session.empty_session() - session.app_id = app_id - return session + return Session( + session_id=f"session-{app_id}", + executor=MagicMock(spec=ThreadPoolExecutor), + reader=MemoryRequestReader(), + writer=RecordingResponseWriter(), + app_id=app_id, + ) def _executor(model: object) -> PluginExecutor: @@ -23,19 +36,19 @@ def _executor(model: object) -> PluginExecutor: return PluginExecutor(config=MagicMock(), registration=registration) -def _llm_data() -> MagicMock: - data = MagicMock() - data.provider = "provider" - data.model_type = "llm" - data.model = "model" - data.credentials = {} - data.prompt_messages = [] - data.model_parameters = {} - data.tools = None - data.stop = None - data.stream = True - data.user_id = "user" - return data +def _llm_data() -> ModelInvokeLLMRequest: + return ModelInvokeLLMRequest.model_validate({ + "provider": "provider", + "model_type": "llm", + "model": "model", + "credentials": {}, + "prompt_messages": [{"role": "user", "content": "hello"}], + "model_parameters": {}, + "tools": None, + "stop": None, + "stream": True, + "user_id": "user", + }) def _text_embedding_data() -> ModelInvokeTextEmbeddingRequest: @@ -99,3 +112,137 @@ def invoke(*_args: object, **_kwargs: object) -> list[str]: user="user", input_type=EmbeddingInputType.QUERY, ) + + +@pytest.mark.parametrize("finish", ["complete", "error", "close"]) +def test_llm_restores_outer_session_when_stream_finishes(finish: str) -> None: + outer, inner = _session("outer"), _session("inner") + seen: list[Session | None] = [] + + def stream() -> Generator[str]: + try: + seen.append(get_current_session()) + yield "first" + if finish == "error": + msg = "provider disconnected" + raise RuntimeError(msg) + seen.append(get_current_session()) + yield "second" + finally: + seen.append(get_current_session()) + + model = MagicMock(spec=LargeLanguageModel) + model.invoke.return_value = stream() + with use_current_session(outer): + result = _executor(model).invoke_llm(inner, _llm_data()) + assert isinstance(result, Generator) + assert get_current_session() is outer + try: + assert next(result) == "first" + if finish == "close": + result.close() + elif finish == "error": + with pytest.raises(RuntimeError, match="provider disconnected"): + next(result) + else: + assert list(result) == ["second"] + assert get_current_session() is outer + finally: + result.close() + + assert seen == [inner] * (3 if finish == "complete" else 2) + assert get_current_session() is None + + +def test_model_creation_error_restores_outer_session() -> None: + outer, inner = _session("outer"), _session("inner") + + def invoke(*_args: object, **_kwargs: object) -> None: + assert get_current_session() is inner + msg = "model rejected credentials" + raise ValueError(msg) + + model = MagicMock(spec=LargeLanguageModel) + model.invoke.side_effect = invoke + with use_current_session(outer): + with pytest.raises(ValueError, match="model rejected credentials"): + _executor(model).invoke_llm(inner, _llm_data()) + assert get_current_session() is outer + assert get_current_session() is None + + +def test_nested_model_invocation_restores_callers_session() -> None: + outer, inner = _session("outer"), _session("inner") + + def embed(*_args: object, **_kwargs: object) -> list[str]: + assert get_current_session() is inner + return ["embedding"] + + embedding = MagicMock(spec=TextEmbeddingModel) + embedding.invoke.side_effect = embed + + def invoke(*_args: object, **_kwargs: object) -> dict[str, str]: + assert get_current_session() is outer + result = _executor(embedding).invoke_text_embedding( + inner, _text_embedding_data() + ) + assert result == ["embedding"] + assert get_current_session() is outer + return {"answer": "done"} + + model = MagicMock(spec=LargeLanguageModel) + model.invoke.side_effect = invoke + assert _executor(model).invoke_llm(outer, _llm_data()) == {"answer": "done"} + assert get_current_session() is None + + +@pytest.mark.parametrize("cancel_first", [False, True]) +def test_concurrent_model_streams_isolate_sessions(cancel_first: bool) -> None: + sessions = [_session("app-a"), _session("app-b")] + entered = [Event(), Event()] + resume = [Event(), Event()] + cleaned: list[Session | None] = [] + + def consume(index: int) -> list[str]: + session = sessions[index] + + def stream() -> Generator[str]: + try: + assert get_current_session() is session + yield f"{session.app_id}-first" + entered[index].set() + assert resume[index].wait(timeout=2), "consumer was not resumed" + assert get_current_session() is session + yield f"{session.app_id}-last" + finally: + cleaned.append(get_current_session()) + + model = MagicMock(spec=LargeLanguageModel) + model.invoke.return_value = stream() + result = _executor(model).invoke_llm(session, _llm_data()) + assert isinstance(result, Generator) + try: + values = list(result) + assert get_current_session() is None + return values + finally: + result.close() + + workers = [gevent.spawn(consume, index) for index in range(2)] + try: + assert all(event.wait(timeout=2) for event in entered) + assert get_current_session() is None + if cancel_first: + workers[0].kill(block=True, timeout=2) + assert workers[0].dead + else: + resume[0].set() + assert workers[0].get(timeout=2) == ["app-a-first", "app-a-last"] + resume[1].set() + assert workers[1].get(timeout=2) == ["app-b-first", "app-b-last"] + assert cleaned == sessions + assert get_current_session() is None + finally: + for event in resume: + event.set() + gevent.killall(workers, block=True, timeout=2) diff --git a/tests/core/test_yaml_loader.py b/tests/core/test_yaml_loader.py index eba7a78d..f8e30975 100644 --- a/tests/core/test_yaml_loader.py +++ b/tests/core/test_yaml_loader.py @@ -5,27 +5,14 @@ from dify_plugin.core.utils.yaml_loader import load_yaml_file -class SwitchingPath: - def __init__(self, paths: list[Path]) -> None: - self.paths = paths - self.calls = 0 +def test_load_yaml_preserves_unicode_and_nested_values(tmp_path: Path) -> None: + file_path = tmp_path / "provider.yaml" + file_path.write_text("label: 你好\nconfig:\n enabled: true\n retries: 0\n") - def __fspath__(self) -> str: - path = self.paths[self.calls] - self.calls += 1 - return str(path) - - -class UnstringableMissingPath: - def __init__(self, file_path: Path, error: RuntimeError) -> None: - self.file_path = file_path - self.error = error - - def __fspath__(self) -> str: - return str(self.file_path) - - def __str__(self) -> str: - raise self.error + assert load_yaml_file(str(file_path)) == { + "label": "你好", + "config": {"enabled": True, "retries": 0}, + } def test_cyclic_symlink_raises_by_default(tmp_path: Path) -> None: @@ -49,24 +36,3 @@ def test_ignore_error_covers_path_probe_errors(tmp_path: Path) -> None: file_path = tmp_path / ("x" * 5000) assert load_yaml_file(str(file_path), ignore_error=True) == {} - - -def test_pathlike_is_resolved_for_exists_and_open(tmp_path: Path) -> None: - first_path = tmp_path / "first.yaml" - second_path = tmp_path / "second.yaml" - first_path.write_text("source: first", encoding="utf-8") - second_path.write_text("source: second", encoding="utf-8") - file_path = SwitchingPath([first_path, second_path]) - - assert load_yaml_file(file_path) == {"source": "second"} - assert file_path.calls == 2 - - -def test_missing_path_preserves_string_conversion_error(tmp_path: Path) -> None: - error = RuntimeError() - file_path = UnstringableMissingPath(tmp_path / "missing.yaml", error) - - with pytest.raises(RuntimeError) as exc_info: - load_yaml_file(file_path) - - assert exc_info.value is error diff --git a/tests/entities/models/test_provider_manifest.py b/tests/entities/models/test_provider_manifest.py new file mode 100644 index 00000000..8d788220 --- /dev/null +++ b/tests/entities/models/test_provider_manifest.py @@ -0,0 +1,80 @@ +from pathlib import Path + +import pytest +import yaml +from pydantic import ValidationError + +from dify_plugin.entities.model import ModelType +from dify_plugin.entities.model.provider import ModelProviderConfiguration +from tests.core.plugin_fixtures import contract_source + +__all__ = ["contract_source"] + + +@pytest.mark.parametrize("existing_position", [False, True]) +def test_model_glob_loads_position_separately_from_models( + contract_source: Path, existing_position: bool +) -> None: + # OpenAI's predefined glob includes _position.yaml alongside model declarations. + data = yaml.safe_load((contract_source / "models/provider.yaml").read_text()) + data["models"]["llm"]["predefined"] = ["models/llm.yaml", "models/_position.yaml"] + (contract_source / "models/_position.yaml").write_text("- echo\n") + if existing_position: + data["position"] = {"tts": ["voice"]} + + provider = ModelProviderConfiguration.model_validate(data) + + assert len(provider.models) == 6 + assert [ + item.model for item in provider.models if item.model_type == ModelType.LLM + ] == ["echo"] + assert provider.position.llm == ["echo"] + assert provider.position.tts == (["voice"] if existing_position else []) + + +@pytest.mark.parametrize("models", [[], "models/*.yaml"]) +def test_model_provider_rejects_invalid_glob_mapping( + contract_source: Path, models: object +) -> None: + data = yaml.safe_load((contract_source / "models/provider.yaml").read_text()) + data["models"] = models + + with pytest.raises((TypeError, ValidationError), match="models"): + ModelProviderConfiguration.model_validate(data) + + +def test_model_provider_reports_empty_matched_file(contract_source: Path) -> None: + data = yaml.safe_load((contract_source / "models/provider.yaml").read_text()) + data["models"]["llm"]["predefined"] = ["models/empty.yaml"] + (contract_source / "models/empty.yaml").write_text("") + + with pytest.raises( + ValidationError, match=r"Error loading model entity:.*empty\.yaml" + ): + ModelProviderConfiguration.model_validate(data) + + +def test_model_provider_manifest_can_be_reused_without_losing_globs( + contract_source: Path, +) -> None: + data = yaml.safe_load((contract_source / "models/provider.yaml").read_text()) + first = ModelProviderConfiguration.model_validate(data) + second = ModelProviderConfiguration.model_validate(data) + + assert first == second + + +def test_model_providers_keep_updates_to_shared_position_data( + contract_source: Path, +) -> None: + data = yaml.safe_load((contract_source / "models/provider.yaml").read_text()) + data["models"]["llm"]["predefined"] = ["models/llm.yaml", "models/_position.yaml"] + data["position"] = {"llm": ["initial"]} + other_data = {**data, "models": {}} + (contract_source / "models/_position.yaml").write_text("- echo\n") + + first = ModelProviderConfiguration.model_validate(data) + second = ModelProviderConfiguration.model_validate(other_data) + + assert first.position.llm == ["echo"] + assert second.position.llm == ["echo"] diff --git a/tests/entities/plugin/test_declaration.py b/tests/entities/plugin/test_declaration.py index a56f0d82..ad04582b 100644 --- a/tests/entities/plugin/test_declaration.py +++ b/tests/entities/plugin/test_declaration.py @@ -1,3 +1,6 @@ +import pytest +from pydantic import ValidationError + from dify_plugin.core.entities.plugin.setup import ( PluginArch, PluginConfiguration, @@ -5,17 +8,18 @@ ) -def test_optional_minimum_dify_required() -> None: +@pytest.mark.parametrize("minimum_version", [None, "1.0.0", "1.0.0rc1"]) +def test_optional_minimum_dify_required(minimum_version: str | None) -> None: meta = PluginConfiguration.Meta( version="1.0.0", arch=[PluginArch.AMD64, PluginArch.ARM64], runner=PluginConfiguration.Meta.PluginRunner( language=PluginLanguage.PYTHON, version="3.10", entrypoint="main.py" ), - minimum_dify_version="1.0.0", + minimum_dify_version=minimum_version, ) - assert meta.minimum_dify_version == "1.0.0" + assert meta.minimum_dify_version == minimum_version meta = PluginConfiguration.Meta( version="1.0.0", @@ -26,3 +30,16 @@ def test_optional_minimum_dify_required() -> None: ) assert meta.minimum_dify_version is None + + +@pytest.mark.parametrize("minimum_version", ["", "not-a-version"]) +def test_invalid_minimum_dify_version_is_rejected(minimum_version: str) -> None: + with pytest.raises(ValidationError, match="Invalid version format"): + PluginConfiguration.Meta( + version="1.0.0", + arch=[PluginArch.AMD64], + runner=PluginConfiguration.Meta.PluginRunner( + language=PluginLanguage.PYTHON, version="3.12", entrypoint="main.py" + ), + minimum_dify_version=minimum_version, + ) diff --git a/tests/entities/test_datasource.py b/tests/entities/test_datasource.py index 2a6d44ad..ab1fde90 100644 --- a/tests/entities/test_datasource.py +++ b/tests/entities/test_datasource.py @@ -1,10 +1,13 @@ import base64 +import pytest + from dify_plugin.entities.datasource import ( DatasourceMessage, OnlineDriveBrowseFilesResponse, OnlineDriveFile, ) +from dify_plugin.entities.datasource_manifest import DatasourceEntity def test_online_drive_file_keeps_legacy_payload_compatible() -> None: @@ -71,3 +74,32 @@ def test_datasource_blob_message_keeps_online_drive_remote_metadata() -> None: assert serialized["meta"] == {"remote_metadata": metadata} assert serialized["message"]["blob"] == base64.b64encode(b"content").decode() + + +@pytest.mark.parametrize(("default", "expected"), [(False, 0.0), (True, 1.0)]) +def test_datasource_manifest_keeps_legacy_numeric_boolean_defaults( + default: bool, + expected: float, +) -> None: + datasource = DatasourceEntity.model_validate({ + "identity": { + "author": "langgenius", + "name": "crawl", + "label": {"en_US": "Crawl"}, + }, + "description": {"en_US": "Crawl a website"}, + "extra": {"python": {"source": "datasources/crawl.py"}}, + "parameters": [ + { + "name": "only_main_content", + "type": "boolean", + "default": default, + "label": {"en_US": "Only main content"}, + "description": {"en_US": "Exclude navigation and footer content"}, + } + ], + }) + + serialized = datasource.model_dump(mode="json")["parameters"][0]["default"] + assert type(serialized) is float + assert serialized == expected diff --git a/tests/entities/test_invoke_message.py b/tests/entities/test_invoke_message.py index e01f1302..35790a50 100644 --- a/tests/entities/test_invoke_message.py +++ b/tests/entities/test_invoke_message.py @@ -1,7 +1,85 @@ +import base64 +from copy import deepcopy + +import pytest +from hypothesis import example, given +from hypothesis import strategies as st +from pydantic import ValidationError + +from dify_plugin.entities.agent import AgentInvokeMessage from dify_plugin.entities.datasource import DatasourceMessage +from dify_plugin.entities.invoke_message import InvokeMessage from dify_plugin.entities.tool import ToolInvokeMessage +@pytest.mark.parametrize( + "message_type", [ToolInvokeMessage, AgentInvokeMessage, DatasourceMessage] +) +@pytest.mark.parametrize("chunked", [False, True]) +@given(blob=st.binary(max_size=32_768)) +@example(blob=b"") +@example(blob=bytes(range(256))) +def test_blob_payload_roundtrips_with_existing_decode_behavior( + message_type: type[InvokeMessage], blob: bytes, chunked: bool +) -> None: + message = {"blob": base64.b64encode(blob).decode()} + if chunked: + message.update(id="file-1", sequence=0, total_length=len(blob), end=True) + payload = { + "type": "blob_chunk" if chunked else "blob", + "message": message, + "meta": { + "mime_type": "application/octet-stream", + "file_name" + if message_type is DatasourceMessage + else "filename": "报告.bin", + }, + } + original = deepcopy(payload) + + first = message_type.model_validate(payload) + restored = message_type.model_validate_json(first.model_dump_json()) + + assert first == restored + assert isinstance( + first.message, InvokeMessage.BlobMessage | InvokeMessage.BlobChunkMessage + ) + assert first.message.blob == blob + assert first.model_dump(mode="json") == original + assert payload["message"]["blob"] == blob + + +@pytest.mark.parametrize( + "message_type", [ToolInvokeMessage, AgentInvokeMessage, DatasourceMessage] +) +def test_blob_decode_preserves_shared_message_metadata( + message_type: type[InvokeMessage], +) -> None: + shared = {"blob": "dGVzdA=="} + + message = message_type.model_validate({ + "type": "blob", + "message": shared, + "meta": shared, + }) + + assert message.message.blob == b"test" + assert message.model_dump(mode="json")["meta"] == {"blob": "test"} + + +@pytest.mark.parametrize("value", [None, 0, False, {}, []]) +def test_stream_variables_reject_non_string_values(value: object) -> None: + with pytest.raises(ValidationError, match="must be a string"): + InvokeMessage.VariableMessage( + variable_name="answer", variable_value=value, stream=True + ) + + message = InvokeMessage.VariableMessage( + variable_name="answer", variable_value=value + ) + assert message.variable_value == value + + def test_tool_invoke_message_accepts_binary_link_type() -> None: message = ToolInvokeMessage.model_validate({ "type": "binary_link", diff --git a/tests/entities/test_manifest_loading.py b/tests/entities/test_manifest_loading.py new file mode 100644 index 00000000..0ca1f78e --- /dev/null +++ b/tests/entities/test_manifest_loading.py @@ -0,0 +1,266 @@ +from copy import deepcopy +from pathlib import Path +from types import MappingProxyType + +import pytest +import yaml +from pydantic import BaseModel, ValidationError + +from dify_plugin.entities.agent import AgentStrategyProviderConfiguration +from dify_plugin.entities.datasource_manifest import DatasourceProviderManifest +from dify_plugin.entities.endpoint import EndpointProviderConfiguration +from dify_plugin.entities.tool import ToolProviderConfiguration +from dify_plugin.entities.trigger import TriggerProviderConfiguration +from tests.core.plugin_fixtures import contract_source + +__all__ = ["contract_source"] + + +@pytest.fixture( + params=[ + pytest.param( + (ToolProviderConfiguration, "provider/contract.yaml", "tools"), id="tool" + ), + pytest.param( + (AgentStrategyProviderConfiguration, "agents/provider.yaml", "strategies"), + id="agent", + ), + pytest.param( + (DatasourceProviderManifest, "datasources/website.yaml", "datasources"), + id="datasource", + ), + pytest.param( + (TriggerProviderConfiguration, "triggers/events.yaml", "events"), + id="trigger", + ), + pytest.param( + (EndpointProviderConfiguration, "endpoints/provider.yaml", "endpoints"), + id="endpoint", + ), + ] +) +def provider_manifest( + request: pytest.FixtureRequest, contract_source: Path +) -> tuple[type[BaseModel], dict, str]: + model, path, field = request.param + return model, yaml.safe_load((contract_source / path).read_text()), field + + +@pytest.mark.parametrize("references", ["child.yaml", [42]]) +def test_provider_rejects_invalid_declaration_references( + provider_manifest: tuple[type[BaseModel], dict, str], references: object +) -> None: + model, data, field = provider_manifest + data[field] = references + + with pytest.raises((TypeError, ValidationError)): + model.model_validate(data) + + +@pytest.mark.parametrize( + ("contents", "cause"), + [ + (None, FileNotFoundError), + ("[", yaml.YAMLError), + ("{}", (KeyError, ValidationError)), + ], + ids=["missing-file", "invalid-yaml", "missing-fields"], +) +def test_provider_reports_failed_declaration_and_keeps_original_cause( + provider_manifest: tuple[type[BaseModel], dict, str], + contract_source: Path, + contents: str | None, + cause: type[Exception] | tuple[type[Exception], ...], +) -> None: + model, data, field = provider_manifest + path = contract_source / "broken.yaml" + if contents is not None: + path.write_text(contents) + data[field] = ["broken.yaml"] + + with pytest.raises( + ValidationError, match=r"Error loading .* configuration" + ) as caught: + model.model_validate(data) + + error = caught.value.errors()[0] + assert error["loc"] == (field,) + assert isinstance(error["ctx"]["error"].__cause__, cause) + if contents is None: + assert "broken.yaml" in str(caught.value) + + +def test_provider_loads_declarations_in_manifest_order( + provider_manifest: tuple[type[BaseModel], dict, str], contract_source: Path +) -> None: + model, data, field = provider_manifest + first = data[field][0] + declaration = ( + deepcopy(first) + if isinstance(first, dict) + else yaml.safe_load((contract_source / first).read_text()) + ) + if field == "endpoints": + declaration["path"] = "/second" + declaration["hidden"] = True + else: + declaration["identity"]["name"] = "second" + declaration["parameters"] = None + (contract_source / "second.yaml").write_text(yaml.safe_dump(declaration)) + data[field] = ["second.yaml", first] + + loaded = getattr(model.model_validate(data), field) + + if field == "endpoints": + assert [endpoint.path for endpoint in loaded] == ["/second", "/echo/"] + assert loaded[0].hidden + else: + assert [item.identity.name for item in loaded] == ["second", "echo"] + assert loaded[0].parameters == [] + assert loaded[0].extra.python.source == declaration["extra"]["python"]["source"] + + +@pytest.mark.parametrize("repair_first", [False, True]) +def test_tool_credentials_survive_revalidation_after_load_or_repair( + contract_source: Path, repair_first: bool +) -> None: + data = yaml.safe_load((contract_source / "provider/contract.yaml").read_text()) + original = deepcopy(data) + if repair_first: + data["tools"] = ["missing.yaml"] + with pytest.raises(ValidationError, match=r"missing\.yaml"): + ToolProviderConfiguration.model_validate(data) + data["tools"] = original["tools"] + else: + ToolProviderConfiguration.model_validate(data) + + reloaded = ToolProviderConfiguration.model_validate(data) + + assert [credential.name for credential in reloaded.credentials_schema] == [ + "api_key" + ] + assert reloaded.tools[0].identity.name == "echo" + + +@pytest.mark.parametrize("immutable_field", [None, "label", "ignored"]) +@pytest.mark.parametrize( + ("model", "path", "field", "result_field"), + [ + ( + ToolProviderConfiguration, + "provider/contract.yaml", + "credentials_for_provider", + "credentials_schema", + ), + ( + TriggerProviderConfiguration, + "triggers/events.yaml", + "credentials_schema", + "subscription_schema", + ), + ], + ids=["tool", "trigger"], +) +def test_provider_credentials_keep_legacy_shared_yaml_mapping_names( + contract_source: Path, + model: type[ToolProviderConfiguration | TriggerProviderConfiguration], + path: str, + field: str, + result_field: str, + immutable_field: str | None, +) -> None: + data = yaml.safe_load((contract_source / path).read_text()) + credential = {"type": "secret-input", "label": {"en_US": "Token"}} + if immutable_field == "label": + credential["label"] = MappingProxyType(credential["label"]) + data[field] = {"access_token": credential, "refresh_token": credential} + if model is TriggerProviderConfiguration: + data["subscription_schema"] = [credential, credential] + if immutable_field == "ignored": + data["unused"] = MappingProxyType({}) + + result = model.model_validate(data) + + assert [item.name for item in getattr(result, result_field)] == [ + "refresh_token", + "refresh_token", + ] + assert model.model_validate(data) == result + + +def test_trigger_shared_yaml_anchor_keeps_legacy_revalidation_values( + contract_source: Path, +) -> None: + data = yaml.safe_load((contract_source / "triggers/events.yaml").read_text()) + data.update( + yaml.safe_load(""" +credentials_schema: + outer: &credential + type: secret-input + label: + en_US: Token +subscription_schema: + - *credential +subscription_constructor: + credentials_schema: + inner: *credential +events: [] +""") + ) + + first = TriggerProviderConfiguration.model_validate(data) + second = TriggerProviderConfiguration.model_validate(data) + + assert [item.name for item in first.subscription_schema] == ["outer"] + assert [item.name for item in second.subscription_schema] == ["inner"] + assert first.subscription_constructor is not None + assert second.subscription_constructor is not None + assert [ + item.name for item in first.subscription_constructor.credentials_schema + ] == ["inner"] + assert [ + item.name for item in second.subscription_constructor.credentials_schema + ] == ["inner"] + + +@pytest.mark.parametrize(("default", "expected"), [(False, 0), (True, 1)]) +def test_agent_manifest_keeps_legacy_numeric_boolean_defaults( + contract_source: Path, default: bool, expected: int +) -> None: + path = contract_source / "agents/echo.yaml" + strategy = yaml.safe_load(path.read_text()) + strategy["parameters"] = [ + { + "name": "enable_search", + "type": "boolean", + "default": default, + "label": {"en_US": "Enable search"}, + } + ] + path.write_text(yaml.safe_dump(strategy)) + provider = yaml.safe_load((contract_source / "agents/provider.yaml").read_text()) + + configuration = AgentStrategyProviderConfiguration.model_validate(provider) + + serialized = configuration.model_dump(mode="json")["strategies"][0]["parameters"][ + 0 + ]["default"] + assert type(serialized) is int + assert serialized == expected + + +def test_endpoint_manifest_can_mix_paths_mappings_and_parsed_declarations( + contract_source: Path, +) -> None: + data = yaml.safe_load((contract_source / "endpoints/provider.yaml").read_text()) + parsed = EndpointProviderConfiguration.model_validate(data).endpoints[0] + (contract_source / "endpoint.yaml").write_text(yaml.safe_dump(data["endpoints"][0])) + data["endpoints"] = ["endpoint.yaml", data["endpoints"][0], parsed] + + result = EndpointProviderConfiguration.model_validate(data) + + assert result.endpoints == [parsed, parsed, parsed] + assert ( + EndpointProviderConfiguration.model_validate_json(result.model_dump_json()) + == result + ) diff --git a/tests/entities/test_trigger.py b/tests/entities/test_trigger.py new file mode 100644 index 00000000..91ba3b90 --- /dev/null +++ b/tests/entities/test_trigger.py @@ -0,0 +1,114 @@ +from types import MappingProxyType + +import pytest +from pydantic import ValidationError + +from dify_plugin.entities.trigger import TriggerSubscriptionConstructorConfiguration + + +@pytest.mark.parametrize("mapping", [False, True]) +def test_subscription_credentials_support_repeated_normalization( + mapping: bool, +) -> None: + # Twilio declares credentials by name; the contract plugin uses a list. + credentials = { + "account_sid": {"type": "text-input", "label": {"en_US": "Account SID"}}, + "auth_token": {"type": "secret-input", "label": {"en_US": "Auth Token"}}, + } + data = { + "credentials_schema": credentials + if mapping + else [{"name": name, **value} for name, value in credentials.items()], + "oauth_schema": { + "client_schema": [ + { + "name": "client_id", + "type": "secret-input", + "label": {"en_US": "Client ID"}, + } + ] + }, + } + first = TriggerSubscriptionConstructorConfiguration.model_validate(data) + second = TriggerSubscriptionConstructorConfiguration.model_validate(data) + + assert first == second + assert [credential.name for credential in first.credentials_schema] == [ + "account_sid", + "auth_token", + ] + assert first.oauth_schema.client_schema[0].name == "client_id" + + +@pytest.mark.parametrize("immutable_field", [None, "label", "ignored"]) +def test_subscription_credentials_keep_legacy_shared_yaml_mapping_names( + immutable_field: str | None, +) -> None: + credential = {"type": "secret-input", "label": {"en_US": "Token"}} + if immutable_field == "label": + credential["label"] = MappingProxyType(credential["label"]) + data = { + "credentials_schema": {"access_token": credential, "refresh_token": credential} + } + if immutable_field == "ignored": + data["unused"] = MappingProxyType({}) + + result = TriggerSubscriptionConstructorConfiguration.model_validate(data) + + assert [item.name for item in result.credentials_schema] == [ + "refresh_token", + "refresh_token", + ] + assert TriggerSubscriptionConstructorConfiguration.model_validate(data) == result + + +@pytest.mark.parametrize("data", [[], "constructor", 42, MappingProxyType({})]) +def test_subscription_constructor_rejects_non_dict_configuration( + data: object, +) -> None: + with pytest.raises((TypeError, ValidationError)): + TriggerSubscriptionConstructorConfiguration.model_validate(data) + + +@pytest.mark.parametrize("credentials", [None, "access_token", 42]) +def test_subscription_constructor_rejects_invalid_credentials_schema( + credentials: object, +) -> None: + with pytest.raises((TypeError, ValidationError), match="credentials_schema"): + TriggerSubscriptionConstructorConfiguration.model_validate({ + "credentials_schema": credentials + }) + + +@pytest.mark.parametrize( + "data", [None, {}, {"credentials_schema": {}, "oauth_schema": {}}] +) +def test_subscription_constructor_accepts_no_authentication(data: object) -> None: + # Discord webhook subscriptions use empty credentials and OAuth schemas. + result = TriggerSubscriptionConstructorConfiguration.model_validate(data) + + assert result.credentials_schema == [] + assert result.parameters == [] + assert result.oauth_schema is None or result.oauth_schema.client_schema == [] + + +@pytest.mark.parametrize(("default", "expected"), [(False, 0), (True, 1)]) +def test_subscription_manifest_keeps_legacy_numeric_boolean_defaults( + default: bool, + expected: int, +) -> None: + # Google Calendar declares this as a boolean in its subscription constructor. + result = TriggerSubscriptionConstructorConfiguration.model_validate({ + "parameters": [ + { + "name": "enrich_event_details", + "type": "boolean", + "default": default, + "label": {"en_US": "Fetch Full Event Details"}, + } + ] + }) + + serialized = result.model_dump(mode="json")["parameters"][0]["default"] + assert type(serialized) is int + assert serialized == expected diff --git a/tests/consts/__init__.py b/tests/file/__init__.py similarity index 100% rename from tests/consts/__init__.py rename to tests/file/__init__.py diff --git a/tests/file/test_file.py b/tests/file/test_file.py new file mode 100644 index 00000000..fa5bd909 --- /dev/null +++ b/tests/file/test_file.py @@ -0,0 +1,136 @@ +import httpx +import pytest +import trustme + +from dify_plugin.config.config import DifyPluginEnv +from dify_plugin.file import file as file_module +from dify_plugin.file.entities import FileType +from dify_plugin.file.file import File +from tests.http_server import HTTPServerFactory, Response + +pytestmark = pytest.mark.integration + + +@pytest.mark.parametrize("content", [b"", b"\x00\xff\x80binary", "你好".encode()]) +def test_download_preserves_bytes_and_caches_empty_files( + http_server: HTTPServerFactory, content: bytes +) -> None: + server = http_server([Response(body=content)]) + path = "/files/document?timestamp=1&nonce=n&sign=a%2Fb%2Bc" + file = File(url=server.url + path, type=FileType.DOCUMENT) + + assert file.blob == content + assert file.blob == content + assert len(server.requests) == 1 + assert server.requests[0].method == "GET" + assert server.requests[0].path == path + + +@pytest.mark.parametrize("status", [401, 403, 404, 429, 500]) +def test_failed_download_can_be_retried_without_caching_the_error( + http_server: HTTPServerFactory, status: int +) -> None: + server = http_server([ + Response(status=status, body=b"failed"), + Response(body=b"recovered"), + ]) + file = File(url=server.url + "/file", type=FileType.IMAGE) + + with pytest.raises(httpx.HTTPStatusError) as exc: + _ = file.blob + assert exc.value.response.status_code == status + assert file.blob == b"recovered" + assert len(server.requests) == 2 + + +@pytest.mark.parametrize("url", ["/files/relative", "ftp://example.com/file"]) +def test_download_explains_unsupported_url_configuration(url: str) -> None: + file = File(url=url, type=FileType.DOCUMENT) + + with pytest.raises(ValueError, match="FILES_URL") as exc: + _ = file.blob + assert isinstance(exc.value.__cause__, httpx.UnsupportedProtocol) + + +def test_download_rejects_a_truncated_body(http_server: HTTPServerFactory) -> None: + server = http_server([Response(headers={"Content-Length": "20"}, body=b"short")]) + file = File(url=server.url + "/file", type=FileType.AUDIO) + + with pytest.raises(httpx.RemoteProtocolError): + _ = file.blob + + +def test_download_preserves_redirect_error_without_following_location( + http_server: HTTPServerFactory, +) -> None: + content = b"\x00\xffredirected" + server = http_server([ + Response(status=302, headers={"Location": "../file?sign=a%2Fb"}), + Response(body=content), + ]) + file = File(url=server.url + "/redirect/start", type=FileType.DOCUMENT) + + with pytest.raises(httpx.HTTPStatusError) as exc: + _ = file.blob + assert exc.value.response.status_code == 302 + assert [request.path for request in server.requests] == ["/redirect/start"] + + +def test_download_honors_configured_read_timeout( + http_server: HTTPServerFactory, monkeypatch: pytest.MonkeyPatch +) -> None: + server = http_server([Response(body=b"late", delay=2)]) + monkeypatch.setattr(file_module, "_plugin_config", DifyPluginEnv(HTTPX_TIMEOUT=1)) + file = File(url=server.url + "/file", type=FileType.VIDEO) + + with pytest.raises(httpx.ReadTimeout): + _ = file.blob + + +def test_download_uses_configured_http_proxy( + http_server: HTTPServerFactory, monkeypatch: pytest.MonkeyPatch +) -> None: + proxy = http_server([Response(body=b"proxied")]) + monkeypatch.setenv("HTTP_PROXY", proxy.url) + url = "http://files.invalid/file?timestamp=1&sign=a%2Fb" + + assert File(url=url, type=FileType.DOCUMENT).blob == b"proxied" + assert len(proxy.requests) == 1 + assert proxy.requests[0].path == url + + +def test_download_rejects_untrusted_tls(http_server: HTTPServerFactory) -> None: + server = http_server([Response(body=b"secure")], tls=True) + file = File( + url=server.url.replace("127.0.0.1", "localhost"), type=FileType.DOCUMENT + ) + + with pytest.raises(httpx.ConnectError, match="CERTIFICATE_VERIFY_FAILED"): + _ = file.blob + assert server.requests == [] + + +@pytest.mark.parametrize("matching_hostname", [True, False]) +def test_download_honors_custom_ca_and_verifies_hostname( + http_server: HTTPServerFactory, + monkeypatch: pytest.MonkeyPatch, + matching_hostname: bool, + tls_ca: trustme.CA, +) -> None: + server = http_server([Response(body=b"secure")], tls=True) + url = ( + server.url.replace("127.0.0.1", "localhost") + if matching_hostname + else server.url + ) + file = File(url=url, type=FileType.DOCUMENT) + + with tls_ca.cert_pem.tempfile() as cert: + monkeypatch.setenv("SSL_CERT_FILE", cert) + if matching_hostname: + assert file.blob == b"secure" + assert len(server.requests) == 1 + else: + with pytest.raises(httpx.ConnectError, match="CERTIFICATE_VERIFY_FAILED"): + _ = file.blob + assert server.requests == [] diff --git a/tests/file/test_upload.py b/tests/file/test_upload.py new file mode 100644 index 00000000..0516cae0 --- /dev/null +++ b/tests/file/test_upload.py @@ -0,0 +1,137 @@ +import json +from email import policy +from email.parser import BytesParser + +import pytest +from pydantic import ValidationError +from pytest_mock import MockerFixture + +from dify_plugin.core.entities.invocation import InvokeType +from dify_plugin.invocations.file import File, UploadFileResponse +from tests.http_server import HTTPServerFactory, Response + +pytestmark = pytest.mark.integration + + +@pytest.mark.parametrize("content", [b"", b"\x00\xff\x80binary", "你好".encode()]) +def test_upload_preserves_signed_url_multipart_and_preview_url( + http_server: HTTPServerFactory, mocker: MockerFixture, content: bytes +) -> None: + upload_data = { + "id": "uploaded-file", + "name": "résumé.png", + "size": len(content), + "extension": "png", + "mime_type": "image/png", + "preview_url": "https://files.example/preview?sign=retained", + } + server = http_server([Response(status=201, body=json.dumps(upload_data).encode())]) + path = "/files/upload/for-plugin?timestamp=1&nonce=n&sign=a%2Fb%2Bc&tenant_id=t" + file = File() + request_url = mocker.patch.object( + file, "_backwards_invoke", return_value=iter([{"url": server.url + path}]) + ) + + result = file.upload("résumé.png", content, "image/png") + + request_url.assert_called_once_with( + InvokeType.UploadFile, dict, {"filename": "résumé.png", "mimetype": "image/png"} + ) + assert result.model_dump(exclude={"type"}) == upload_data + assert result.to_app_parameter() == { + "tool_file_id": "uploaded-file", + "transfer_method": "tool_file", + "type": "image", + } + assert len(server.requests) == 1 + request = server.requests[0] + assert (request.method, request.path) == ("POST", path) + message = BytesParser(policy=policy.default).parsebytes( + f"Content-Type: {request.headers['Content-Type']}\r\n\r\n".encode() + + request.body + ) + parts = list(message.iter_parts()) + assert len(parts) == 1 + assert parts[0].get_param("name", header="content-disposition") == "file" + assert parts[0].get_filename() == "résumé.png" + assert parts[0].get_content_type() == "image/png" + assert parts[0].get_payload(decode=True) == content + + +@pytest.mark.parametrize("status", [200, 400, 403, 413, 415, 500]) +def test_upload_surfaces_rejected_or_unexpected_status( + http_server: HTTPServerFactory, mocker: MockerFixture, status: int +) -> None: + server = http_server([Response(status=status, body=b"rejected upload")]) + file = File() + mocker.patch.object( + file, "_backwards_invoke", return_value=iter([{"url": server.url + "/upload"}]) + ) + + with pytest.raises(Exception, match=f"status code: {status}.*rejected upload"): + file.upload("file.txt", b"test", "text/plain") + assert len(server.requests) == 1 + + +@pytest.mark.parametrize( + ("responses", "message"), + [([], "empty response"), ([{}], "signed url"), ([{"url": ""}], "signed url")], +) +def test_upload_requires_a_signed_url( + mocker: MockerFixture, responses: list[dict], message: str +) -> None: + file = File() + mocker.patch.object(file, "_backwards_invoke", return_value=iter(responses)) + + with pytest.raises(Exception, match=message): + file.upload("file.txt", b"test", "text/plain") + + +@pytest.mark.parametrize( + ("mime_type", "expected"), + [ + ("image/png", "image"), + ("video/mp4", "video"), + ("audio/wav", "audio"), + ("application/pdf", "document"), + ("application/octet-stream", "document"), + ], +) +def test_upload_result_is_usable_as_an_app_file(mime_type: str, expected: str) -> None: + result = UploadFileResponse( + id="file-id", name="file", size=0, extension="", mime_type=mime_type + ) + + assert result.type == expected + assert result.to_app_parameter() == { + "tool_file_id": "file-id", + "transfer_method": "tool_file", + "type": expected, + } + + +def test_upload_accepts_an_explicit_response_type() -> None: + result = UploadFileResponse( + id="file-id", + name="file", + size=0, + extension="", + mime_type="image/png", + type="image", + ) + + assert result.type == UploadFileResponse.Type.IMAGE + + +@pytest.mark.parametrize("body", [b"not JSON", b'{"id": "incomplete"}']) +def test_upload_rejects_invalid_success_payload( + http_server: HTTPServerFactory, mocker: MockerFixture, body: bytes +) -> None: + server = http_server([Response(status=201, body=body)]) + file = File() + mocker.patch.object( + file, "_backwards_invoke", return_value=iter([{"url": server.url + "/upload"}]) + ) + + with pytest.raises((ValueError, ValidationError)): + file.upload("file.txt", b"test", "text/plain") diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugin_cli.py b/tests/fixtures/plugin_cli.py new file mode 100644 index 00000000..1c843ce6 --- /dev/null +++ b/tests/fixtures/plugin_cli.py @@ -0,0 +1,99 @@ +"""Emulate the CLI JSON boundary for PluginRunner process tests.""" + +import json +import signal +import sys +import time +from pathlib import Path +from types import FrameType +from typing import Any, TextIO + + +def emit( + event: str, + invoke_id: str = "", + response: dict[str, Any] | None = None, + output: TextIO = sys.stdout, +) -> None: + output.write( + json.dumps({"invoke_id": invoke_id, "type": event, "response": response or {}}) + + "\n" + ) + output.flush() + + +def respond(request: dict[str, Any]) -> None: + emit("plugin_response", request["invoke_id"], request["request"]) + emit("plugin_invoke_end", request["invoke_id"], {"info": "plugin invoke end"}) + + +def main() -> None: + if sys.argv[1:] == ["version"]: + sys.stdout.write("0.1.0\n") + return + if sys.argv[1:3] == ["plugin", "package"]: + if Path(sys.argv[3]).name == "fail-package": + sys.stderr.write("package failed\n") + raise SystemExit(1) + Path(sys.argv[5]).write_bytes(b"local fixture package") + return + + scenario = sys.argv[7] if sys.argv[6:7] == ["--scenario"] else "echo" + if scenario == "exit-before-ready": + return + if scenario == "error-before-ready": + emit("error", response={"error": "startup failed"}) + return + if scenario == "graceful-termination": + gate = Path(sys.argv[8]) + + def shutdown(_signal: int, _frame: FrameType | None) -> None: + gate.with_suffix(".started").touch() + while not gate.exists(): + time.sleep(0.01) + gate.with_suffix(".completed").touch() + raise SystemExit + + signal.signal(signal.SIGTERM, shutdown) + if scenario == "never-ready": + sys.stdin.read() + return + emit("plugin_ready", response={"info": "plugin loaded"}) + if scenario in {"stderr", "stderr-error"}: + event = "error" if scenario == "stderr-error" else "info" + emit(event, response={"info": "x" * 200_000}, output=sys.stderr) + + pending = [] + for line in sys.stdin: + request = json.loads(line) + if scenario == "exit-on-invoke": + return + if scenario == "error-on-invoke": + # The CLI emits no end frame when invocation setup fails. + emit("error", request["invoke_id"], {"error": "invocation failed"}) + elif scenario == "record-request": + emit("plugin_response", request["invoke_id"], request) + emit( + "plugin_invoke_end", request["invoke_id"], {"info": "plugin invoke end"} + ) + elif scenario == "invalid-response": + emit("info", request["invoke_id"]) + elif scenario == "mixed": + sys.stdout.write("\nnot json\n") + emit("info", response={"info": "fixture log"}) + emit("plugin_response", "unknown-invocation", {"value": "ignored"}) + respond(request) + elif scenario == "no-response": + continue + elif scenario == "reverse": + pending.append(request) + if len(pending) == 2: + for item in reversed(pending): + respond(item) + pending.clear() + else: + respond(request) + + +if __name__ == "__main__": + main() diff --git a/tests/fixtures/plugins/__init__.py b/tests/fixtures/plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/__init__.py b/tests/fixtures/plugins/contract_plugin/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/_assets/icon.svg b/tests/fixtures/plugins/contract_plugin/_assets/icon.svg new file mode 100644 index 00000000..934c0dec --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/_assets/icon.svg @@ -0,0 +1 @@ + diff --git a/tests/fixtures/plugins/contract_plugin/agents/__init__.py b/tests/fixtures/plugins/contract_plugin/agents/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/agents/echo.py b/tests/fixtures/plugins/contract_plugin/agents/echo.py new file mode 100644 index 00000000..9d18f1cd --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/agents/echo.py @@ -0,0 +1,9 @@ +from collections.abc import Generator + +from dify_plugin.entities.agent import AgentInvokeMessage +from dify_plugin.interfaces.agent import AgentStrategy + + +class EchoAgent(AgentStrategy): + def _invoke(self, parameters: dict) -> Generator[AgentInvokeMessage, None, None]: + yield self.create_text_message(parameters["message"]) diff --git a/tests/fixtures/plugins/contract_plugin/agents/echo.yaml b/tests/fixtures/plugins/contract_plugin/agents/echo.yaml new file mode 100644 index 00000000..76620e69 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/agents/echo.yaml @@ -0,0 +1,11 @@ +identity: + author: sdk_tests + name: echo + label: + en_US: Echo +description: + en_US: Return the requested message. +parameters: [] +extra: + python: + source: agents/echo.py diff --git a/tests/fixtures/plugins/contract_plugin/agents/provider.yaml b/tests/fixtures/plugins/contract_plugin/agents/provider.yaml new file mode 100644 index 00000000..f9ccace6 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/agents/provider.yaml @@ -0,0 +1,10 @@ +identity: + author: sdk_tests + name: contract_agents + label: + en_US: Contract agents + description: + en_US: Local SDK test agent. + icon: icon.svg +strategies: + - agents/echo.yaml diff --git a/tests/fixtures/plugins/contract_plugin/datasources/__init__.py b/tests/fixtures/plugins/contract_plugin/datasources/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/datasources/documents.py b/tests/fixtures/plugins/contract_plugin/datasources/documents.py new file mode 100644 index 00000000..c42b4f42 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/documents.py @@ -0,0 +1,51 @@ +from collections.abc import Generator, Mapping +from typing import Any + +from dify_plugin.entities.datasource import ( + DatasourceGetPagesResponse, + DatasourceMessage, + GetOnlineDocumentPageContentRequest, + OnlineDocumentInfo, + OnlineDocumentPage, +) +from dify_plugin.interfaces.datasource.online_document import OnlineDocumentDatasource + + +class ContractDocuments(OnlineDocumentDatasource): + def _get_pages( + self, datasource_parameters: Mapping[str, Any] + ) -> DatasourceGetPagesResponse: + if datasource_parameters.get("empty"): + return DatasourceGetPagesResponse(result=[]) + return DatasourceGetPagesResponse( + result=[ + OnlineDocumentInfo( + workspace_id="workspace", + workspace_name="Documents", + workspace_icon="", + total=1, + pages=[ + OnlineDocumentPage( + page_id="page", + page_name="Example", + type="page", + last_edited_time="2026-01-01T00:00:00Z", + ) + ], + ) + ] + ) + + def _get_content( + self, page: GetOnlineDocumentPageContentRequest + ) -> Generator[DatasourceMessage]: + self.finished = False + try: + if page.page_id == "missing": + msg = "page not found" + raise ValueError(msg) + yield self.create_variable_message("content", "page content") + yield self.create_variable_message("page_id", page.page_id) + yield self.create_variable_message("workspace_id", page.workspace_id) + finally: + self.finished = True diff --git a/tests/fixtures/plugins/contract_plugin/datasources/documents.yaml b/tests/fixtures/plugins/contract_plugin/datasources/documents.yaml new file mode 100644 index 00000000..23a6ac1b --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/documents.yaml @@ -0,0 +1,30 @@ +identity: + author: sdk_tests + name: documents + label: + en_US: documents + description: + en_US: Local datasource provider. + icon: icon.svg +provider_type: online_document +credentials_schema: + - name: api_key + type: secret-input + label: + en_US: API key +oauth_schema: + client_schema: + - name: client_id + type: secret-input + label: + en_US: Client ID + credentials_schema: + - name: access_token + type: secret-input + label: + en_US: Access token +datasources: + - datasources/documents_echo.yaml +extra: + python: + source: datasources/provider.py diff --git a/tests/fixtures/plugins/contract_plugin/datasources/documents_echo.yaml b/tests/fixtures/plugins/contract_plugin/datasources/documents_echo.yaml new file mode 100644 index 00000000..08761d1e --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/documents_echo.yaml @@ -0,0 +1,10 @@ +identity: + author: sdk_tests + name: echo + label: + en_US: Echo +description: + en_US: Local documents datasource. +extra: + python: + source: datasources/documents.py diff --git a/tests/fixtures/plugins/contract_plugin/datasources/drive.py b/tests/fixtures/plugins/contract_plugin/datasources/drive.py new file mode 100644 index 00000000..9a0f23f1 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/drive.py @@ -0,0 +1,52 @@ +from collections.abc import Generator + +from dify_plugin.entities.datasource import ( + DatasourceMessage, + OnlineDriveBrowseFilesRequest, + OnlineDriveBrowseFilesResponse, + OnlineDriveDownloadFileRequest, + OnlineDriveFile, + OnlineDriveFileBucket, +) +from dify_plugin.interfaces.datasource.online_drive import OnlineDriveDatasource + + +class ContractDrive(OnlineDriveDatasource): + def _browse_files( + self, request: OnlineDriveBrowseFilesRequest + ) -> OnlineDriveBrowseFilesResponse: + if request.prefix == "empty": + return OnlineDriveBrowseFilesResponse(result=[]) + last_page = bool(request.next_page_parameters) + return OnlineDriveBrowseFilesResponse( + result=[ + OnlineDriveFileBucket( + bucket=request.bucket, + files=[ + OnlineDriveFile( + id="last-file" if last_page else "file", + name="example.txt", + size=4, + type="file", + remote_metadata={"etag": '"v1"'}, + ) + ], + is_truncated=not last_page, + next_page_parameters=None if last_page else {"cursor": "next"}, + ) + ] + ) + + def _download_file( + self, request: OnlineDriveDownloadFileRequest + ) -> Generator[DatasourceMessage]: + self.finished = False + try: + yield self.create_blob_message( + b"data", meta={"mime_type": "text/plain", "file_name": "example.txt"} + ) + if request.id == "broken": + msg = "download interrupted" + raise ValueError(msg) + finally: + self.finished = True diff --git a/tests/fixtures/plugins/contract_plugin/datasources/drive.yaml b/tests/fixtures/plugins/contract_plugin/datasources/drive.yaml new file mode 100644 index 00000000..626e87cc --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/drive.yaml @@ -0,0 +1,30 @@ +identity: + author: sdk_tests + name: drive + label: + en_US: drive + description: + en_US: Local datasource provider. + icon: icon.svg +provider_type: online_drive +credentials_schema: + - name: api_key + type: secret-input + label: + en_US: API key +oauth_schema: + client_schema: + - name: client_id + type: secret-input + label: + en_US: Client ID + credentials_schema: + - name: access_token + type: secret-input + label: + en_US: Access token +datasources: + - datasources/drive_echo.yaml +extra: + python: + source: datasources/provider.py diff --git a/tests/fixtures/plugins/contract_plugin/datasources/drive_echo.yaml b/tests/fixtures/plugins/contract_plugin/datasources/drive_echo.yaml new file mode 100644 index 00000000..edb4571b --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/drive_echo.yaml @@ -0,0 +1,10 @@ +identity: + author: sdk_tests + name: echo + label: + en_US: Echo +description: + en_US: Local drive datasource. +extra: + python: + source: datasources/drive.py diff --git a/tests/fixtures/plugins/contract_plugin/datasources/provider.py b/tests/fixtures/plugins/contract_plugin/datasources/provider.py new file mode 100644 index 00000000..414cfde7 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/provider.py @@ -0,0 +1,49 @@ +from collections.abc import Mapping +from typing import Any +from urllib.parse import urlencode + +from werkzeug import Request + +from dify_plugin.entities.datasource import DatasourceOAuthCredentials +from dify_plugin.interfaces.datasource import DatasourceProvider + + +class ContractDatasourceProvider(DatasourceProvider): + def _validate_credentials(self, credentials: Mapping[str, Any]) -> None: + if credentials.get("api_key") != "valid": + msg = "invalid credentials" + raise ValueError(msg) + + def _oauth_get_authorization_url( + self, redirect_uri: str, system_credentials: Mapping[str, Any] + ) -> str: + return "https://provider.invalid/authorize?" + urlencode({ + "redirect_uri": redirect_uri, + "client_id": system_credentials["client_id"], + }) + + def _oauth_get_credentials( + self, redirect_uri: str, system_credentials: Mapping[str, Any], request: Request + ) -> DatasourceOAuthCredentials: + del redirect_uri, system_credentials + return DatasourceOAuthCredentials( + credentials={ + "access_token": "token-" + request.args["code"], + "refresh_token": "refresh", + }, + name="Contract account", + avatar_url="https://provider.invalid/avatar.png", + ) + + def _oauth_refresh_credentials( + self, + redirect_uri: str, + system_credentials: Mapping[str, Any], + credentials: Mapping[str, Any], + ) -> DatasourceOAuthCredentials: + del redirect_uri, system_credentials + return DatasourceOAuthCredentials( + credentials={**credentials, "access_token": "refreshed"}, + name="Contract account", + avatar_url="https://provider.invalid/avatar.png", + ) diff --git a/tests/fixtures/plugins/contract_plugin/datasources/website.py b/tests/fixtures/plugins/contract_plugin/datasources/website.py new file mode 100644 index 00000000..7b8116c5 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/website.py @@ -0,0 +1,40 @@ +from collections.abc import Generator, Mapping +from typing import Any + +from dify_plugin.entities.datasource import ( + WebsiteCrawlMessage, + WebSiteInfo, + WebSiteInfoDetail, +) +from dify_plugin.interfaces.datasource.website import WebsiteCrawlDatasource + + +class ContractWebsite(WebsiteCrawlDatasource): + def _get_website_crawl( + self, datasource_parameters: Mapping[str, Any] + ) -> Generator[WebsiteCrawlMessage]: + self.finished = False + try: + yield self.create_crawl_message( + WebSiteInfo(status="processing", total=1, completed=0) + ) + if datasource_parameters.get("fail"): + msg = "crawl failed" + raise ValueError(msg) + yield self.create_crawl_message( + WebSiteInfo( + status="completed", + total=1, + completed=1, + web_info_list=[ + WebSiteInfoDetail( + source_url=datasource_parameters["url"], + content="page content", + title="Example", + description="Local test page", + ) + ], + ) + ) + finally: + self.finished = True diff --git a/tests/fixtures/plugins/contract_plugin/datasources/website.yaml b/tests/fixtures/plugins/contract_plugin/datasources/website.yaml new file mode 100644 index 00000000..94c96cdc --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/website.yaml @@ -0,0 +1,30 @@ +identity: + author: sdk_tests + name: website + label: + en_US: website + description: + en_US: Local datasource provider. + icon: icon.svg +provider_type: website_crawl +credentials_schema: + - name: api_key + type: secret-input + label: + en_US: API key +oauth_schema: + client_schema: + - name: client_id + type: secret-input + label: + en_US: Client ID + credentials_schema: + - name: access_token + type: secret-input + label: + en_US: Access token +datasources: + - datasources/website_echo.yaml +extra: + python: + source: datasources/provider.py diff --git a/tests/fixtures/plugins/contract_plugin/datasources/website_echo.yaml b/tests/fixtures/plugins/contract_plugin/datasources/website_echo.yaml new file mode 100644 index 00000000..d5bf5c6c --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/datasources/website_echo.yaml @@ -0,0 +1,10 @@ +identity: + author: sdk_tests + name: echo + label: + en_US: Echo +description: + en_US: Local website datasource. +extra: + python: + source: datasources/website.py diff --git a/tests/fixtures/plugins/contract_plugin/endpoints/__init__.py b/tests/fixtures/plugins/contract_plugin/endpoints/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/endpoints/echo.py b/tests/fixtures/plugins/contract_plugin/endpoints/echo.py new file mode 100644 index 00000000..c0eb6b0e --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/endpoints/echo.py @@ -0,0 +1,15 @@ +from collections.abc import Mapping + +from werkzeug import Request, Response + +from dify_plugin.interfaces.endpoint import Endpoint + + +class EchoEndpoint(Endpoint): + def _invoke(self, r: Request, values: Mapping, settings: Mapping) -> Response: + return Response( + r.get_data(), + status=201, + headers={"X-Item": values["item"], "X-Setting": settings["label"]}, + mimetype="application/octet-stream", + ) diff --git a/tests/fixtures/plugins/contract_plugin/endpoints/provider.yaml b/tests/fixtures/plugins/contract_plugin/endpoints/provider.yaml new file mode 100644 index 00000000..65843dbf --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/endpoints/provider.yaml @@ -0,0 +1,6 @@ +endpoints: + - path: /echo/ + method: POST + extra: + python: + source: endpoints/echo.py diff --git a/tests/fixtures/plugins/contract_plugin/main.py b/tests/fixtures/plugins/contract_plugin/main.py new file mode 100644 index 00000000..06c3d8f0 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/main.py @@ -0,0 +1,24 @@ +import json +import sys +from importlib.metadata import version + +from coverage import Coverage + +import dify_plugin +from dify_plugin import DifyPluginEnv, Plugin + +if __name__ == "__main__": + sys.stderr.write( + json.dumps({ + "sdk_file": dify_plugin.__file__, + "sdk_version": version("dify-plugin"), + "coverage_active": Coverage.current() is not None, + }) + + "\n" + ) + sys.stderr.flush() + Plugin( + DifyPluginEnv( + INSTALL_METHOD="local", MAX_WORKER=4, HEARTBEAT_INTERVAL=60, _env_file=None + ) + ).run() diff --git a/tests/fixtures/plugins/contract_plugin/manifest.yaml b/tests/fixtures/plugins/contract_plugin/manifest.yaml new file mode 100644 index 00000000..d8e88630 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/manifest.yaml @@ -0,0 +1,36 @@ +version: 0.0.1 +type: plugin +author: sdk_tests +name: contract_plugin +label: + en_US: SDK contract fixture +description: + en_US: Local plugin used by SDK tests. +icon: icon.svg +created_at: "2026-01-01T00:00:00Z" +resource: + memory: 1048576 +plugins: + tools: + - provider/contract.yaml + models: + - models/provider.yaml + agent_strategies: + - agents/provider.yaml + endpoints: + - endpoints/provider.yaml + datasources: + - datasources/website.yaml + - datasources/documents.yaml + - datasources/drive.yaml + triggers: + - triggers/events.yaml +meta: + version: 0.0.2 + arch: + - amd64 + - arm64 + runner: + language: python + version: "3.12" + entrypoint: main diff --git a/tests/fixtures/plugins/contract_plugin/models/__init__.py b/tests/fixtures/plugins/contract_plugin/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/models/embedding.yaml b/tests/fixtures/plugins/contract_plugin/models/embedding.yaml new file mode 100644 index 00000000..94fd301a --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/embedding.yaml @@ -0,0 +1,5 @@ +model: echo +model_type: text-embedding +model_properties: + context_size: 1024 + max_chunks: 8 diff --git a/tests/fixtures/plugins/contract_plugin/models/implementations.py b/tests/fixtures/plugins/contract_plugin/models/implementations.py new file mode 100644 index 00000000..90420609 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/implementations.py @@ -0,0 +1,236 @@ +from collections.abc import Generator, Mapping, Sequence +from typing import IO, Literal + +from pydantic import JsonValue + +import dify_plugin as sdk +from dify_plugin.entities.model import EmbeddingInputType +from dify_plugin.entities.model.llm import ( + LLMPollingResult, + LLMPollingStatus, + LLMResult, + LLMResultChunk, + LLMUsage, +) +from dify_plugin.entities.model.message import ( + AssistantPromptMessage, + PromptMessage, + PromptMessageTool, +) +from dify_plugin.entities.model.rerank import ( + MultiModalRerankResult, + RerankDocument, + RerankResult, +) +from dify_plugin.entities.model.text_embedding import ( + EmbeddingUsage, + MultiModalContent, + MultiModalEmbeddingResult, + TextEmbeddingResult, +) +from dify_plugin.errors.model import InvokeError + + +class ContractCredentials: + def validate_credentials(self, model: str, credentials: Mapping) -> None: + del model + if credentials.get("api_key") != "valid": + message = "invalid credentials" + raise ValueError(message) + + @property + def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]: + return {} + + +class ContractLLM(ContractCredentials, sdk.LargeLanguageModel): + def _invoke( + self, + model: str, + credentials: dict, + prompt_messages: list[PromptMessage], + model_parameters: dict, + tools: list[PromptMessageTool] | None = None, + stop: list[str] | None = None, + stream: bool = True, + user: str | None = None, + ) -> LLMResult | Generator[LLMResultChunk, None, None]: + del credentials, model_parameters, tools, stop, user + result = LLMResult( + model=model, + message=AssistantPromptMessage(content=prompt_messages[0].content), + usage=LLMUsage.empty_usage(), + ) + if not stream: + return result + + def chunks() -> Generator[LLMResultChunk, None, None]: + yield result.to_llm_result_chunk() + + return chunks() + + def get_num_tokens( + self, + model: str, + credentials: dict, + prompt_messages: list[PromptMessage], + tools: list[PromptMessageTool] | None = None, + ) -> int: + del model, credentials, tools + return sum(len(message.content) for message in prompt_messages) + + def _start_polling( + self, + model: str, + credentials: dict, + prompt_messages: list[PromptMessage], + model_parameters: dict, + tools: list[PromptMessageTool] | None = None, + stop: list[str] | None = None, + stream: Literal[False] = False, + user: str | None = None, + json_schema: dict[str, JsonValue] | None = None, + ) -> LLMPollingResult: + del credentials, model_parameters, tools, stop, stream, user + return LLMPollingResult( + status=LLMPollingStatus.RUNNING, + plugin_state={ + "model": model, + "content": prompt_messages[0].content, + "json_schema": json_schema, + }, + next_check_after_seconds=1, + ) + + def _check_polling( + self, + model: str, + credentials: dict, + plugin_state: dict[str, JsonValue], + user: str | None = None, + ) -> LLMPollingResult: + del credentials, user + return LLMPollingResult( + status=LLMPollingStatus.SUCCEEDED, + result=LLMResult( + model=model, + message=AssistantPromptMessage(content=plugin_state["content"]), + usage=LLMUsage.empty_usage(), + ), + ) + + +class ContractEmbedding(ContractCredentials, sdk.TextEmbeddingModel): + def _invoke( + self, + model: str, + credentials: dict, + texts: list[str], + user: str | None = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, + ) -> TextEmbeddingResult: + del credentials, user, input_type + return TextEmbeddingResult( + model=model, + embeddings=[[float(len(text))] for text in texts], + usage=EmbeddingUsage( + tokens=0, + total_tokens=0, + unit_price=0, + price_unit=0, + total_price=0, + currency="USD", + latency=0, + ), + ) + + def _invoke_multimodal( + self, + model: str, + credentials: dict, + documents: list[MultiModalContent], + user: str | None = None, + input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT, + ) -> MultiModalEmbeddingResult: + result = self._invoke( + model, credentials, [doc.content for doc in documents], user, input_type + ) + return MultiModalEmbeddingResult(**result.model_dump()) + + def get_num_tokens( + self, model: str, credentials: dict, texts: list[str] + ) -> list[int]: + del model, credentials + return [len(text) for text in texts] + + +class ContractRerank(ContractCredentials, sdk.RerankModel): + def _invoke( + self, + model: str, + credentials: dict, + query: str, + docs: list[str], + score_threshold: float | None = None, + top_n: int | None = None, + user: str | None = None, + ) -> RerankResult: + del credentials, query, score_threshold, top_n, user + return RerankResult( + model=model, docs=[RerankDocument(index=0, text=docs[0], score=0.9)] + ) + + def _invoke_multimodal( + self, + model: str, + credentials: dict, + query: MultiModalContent, + docs: Sequence[MultiModalContent], + score_threshold: float | None = None, + top_n: int | None = None, + user: str | None = None, + ) -> MultiModalRerankResult: + result = self._invoke( + model, + credentials, + query.content, + [doc.content for doc in docs], + score_threshold, + top_n, + user, + ) + return MultiModalRerankResult(**result.model_dump()) + + +class ContractTTS(ContractCredentials, sdk.TTSModel): + def _invoke( + self, + model: str, + tenant_id: str, + credentials: dict, + content_text: str, + voice: str, + user: str | None = None, + ) -> bytes: + del model, tenant_id, credentials, voice, user + return content_text.encode() + + +class ContractSpeechToText(ContractCredentials, sdk.Speech2TextModel): + def _invoke( + self, + model: str, + credentials: dict, + file: IO[bytes], + user: str | None = None, + ) -> str: + del model, credentials, user + return file.read().decode() + + +class ContractModeration(ContractCredentials, sdk.ModerationModel): + def _invoke( + self, model: str, credentials: dict, text: str, user: str | None = None + ) -> bool: + del model, credentials, user + return text == "blocked" diff --git a/tests/fixtures/plugins/contract_plugin/models/llm.yaml b/tests/fixtures/plugins/contract_plugin/models/llm.yaml new file mode 100644 index 00000000..cdf816be --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/llm.yaml @@ -0,0 +1,8 @@ +model: echo +model_type: llm +features: [polling] +model_properties: + mode: chat +parameter_rules: + - name: temperature + use_template: temperature diff --git a/tests/fixtures/plugins/contract_plugin/models/moderation.yaml b/tests/fixtures/plugins/contract_plugin/models/moderation.yaml new file mode 100644 index 00000000..4c5977a9 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/moderation.yaml @@ -0,0 +1,3 @@ +model: echo +model_type: moderation +model_properties: {} diff --git a/tests/fixtures/plugins/contract_plugin/models/provider.py b/tests/fixtures/plugins/contract_plugin/models/provider.py new file mode 100644 index 00000000..ae14fc52 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/provider.py @@ -0,0 +1,8 @@ +from dify_plugin.interfaces.model import ModelProvider + + +class ContractModelProvider(ModelProvider): + def validate_provider_credentials(self, credentials: dict) -> None: + if credentials.get("api_key") != "valid": + message = "invalid credentials" + raise ValueError(message) diff --git a/tests/fixtures/plugins/contract_plugin/models/provider.yaml b/tests/fixtures/plugins/contract_plugin/models/provider.yaml new file mode 100644 index 00000000..8ff07aa1 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/provider.yaml @@ -0,0 +1,29 @@ +provider: contract_models +label: + en_US: Contract models +supported_model_types: + - llm + - text-embedding + - rerank + - tts + - speech2text + - moderation +configurate_methods: + - predefined-model +models: + llm: + predefined: [models/llm.yaml] + text_embedding: + predefined: [models/embedding.yaml] + rerank: + predefined: [models/rerank.yaml] + tts: + predefined: [models/tts.yaml] + speech2text: + predefined: [models/speech2text.yaml] + moderation: + predefined: [models/moderation.yaml] +extra: + python: + provider_source: models/provider.py + model_sources: [models/implementations.py] diff --git a/tests/fixtures/plugins/contract_plugin/models/rerank.yaml b/tests/fixtures/plugins/contract_plugin/models/rerank.yaml new file mode 100644 index 00000000..331fdc26 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/rerank.yaml @@ -0,0 +1,3 @@ +model: echo +model_type: rerank +model_properties: {} diff --git a/tests/fixtures/plugins/contract_plugin/models/speech2text.yaml b/tests/fixtures/plugins/contract_plugin/models/speech2text.yaml new file mode 100644 index 00000000..eca936e9 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/speech2text.yaml @@ -0,0 +1,3 @@ +model: echo +model_type: speech2text +model_properties: {} diff --git a/tests/fixtures/plugins/contract_plugin/models/tts.yaml b/tests/fixtures/plugins/contract_plugin/models/tts.yaml new file mode 100644 index 00000000..6e75c8ef --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/models/tts.yaml @@ -0,0 +1,8 @@ +model: echo +model_type: tts +model_properties: + audio_type: wav + voices: + - name: Echo + mode: echo + language: [en_US] diff --git a/tests/fixtures/plugins/contract_plugin/provider/__init__.py b/tests/fixtures/plugins/contract_plugin/provider/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/provider/contract.py b/tests/fixtures/plugins/contract_plugin/provider/contract.py new file mode 100644 index 00000000..e8cd134b --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/provider/contract.py @@ -0,0 +1,8 @@ +from dify_plugin.interfaces.tool import ToolProvider + + +class ContractProvider(ToolProvider): + def _validate_credentials(self, credentials: dict) -> None: + if credentials.get("api_key") != "valid": + msg = "invalid credentials" + raise ValueError(msg) diff --git a/tests/fixtures/plugins/contract_plugin/provider/contract.yaml b/tests/fixtures/plugins/contract_plugin/provider/contract.yaml new file mode 100644 index 00000000..d348f337 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/provider/contract.yaml @@ -0,0 +1,19 @@ +identity: + author: sdk_tests + name: contract + label: + en_US: Contract + description: + en_US: Local SDK test provider. + icon: icon.svg +credentials_for_provider: + api_key: + type: secret-input + required: true + label: + en_US: API key +tools: + - tools/echo.yaml +extra: + python: + source: provider/contract.py diff --git a/tests/fixtures/plugins/contract_plugin/tools/__init__.py b/tests/fixtures/plugins/contract_plugin/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/tools/echo.py b/tests/fixtures/plugins/contract_plugin/tools/echo.py new file mode 100644 index 00000000..cc2f3b94 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/tools/echo.py @@ -0,0 +1,39 @@ +from collections.abc import Generator +from threading import Barrier + +from dify_plugin.entities import I18nObject +from dify_plugin.entities.tool import ToolInvokeMessage, ToolParameter +from dify_plugin.interfaces.tool import Tool + +peers = Barrier(2) + + +class EchoTool(Tool): + def _invoke(self, tool_parameters: dict) -> Generator[ToolInvokeMessage]: + if tool_parameters.get("wait_for_peer"): + peers.wait(timeout=2) + message = tool_parameters["message"] + yield self.create_text_message(message) + yield self.create_json_message({"echo": message}) + + def _get_runtime_parameters(self) -> list[ToolParameter]: + return [ + ToolParameter( + name="message", + type=ToolParameter.ToolParameterType.STRING, + form=ToolParameter.ToolParameterForm.LLM, + label=I18nObject(en_us="Message"), + human_description=I18nObject(en_us="Message to return."), + llm_description="Message to return.", + required=True, + ), + ToolParameter( + name="wait_for_peer", + type=ToolParameter.ToolParameterType.BOOLEAN, + form=ToolParameter.ToolParameterForm.LLM, + label=I18nObject(en_us="Wait for another invocation"), + human_description=I18nObject(en_us="Synchronize test invocations."), + llm_description="Synchronize test invocations.", + required=False, + ), + ] diff --git a/tests/fixtures/plugins/contract_plugin/tools/echo.yaml b/tests/fixtures/plugins/contract_plugin/tools/echo.yaml new file mode 100644 index 00000000..708ad1ef --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/tools/echo.yaml @@ -0,0 +1,31 @@ +identity: + author: sdk_tests + name: echo + label: + en_US: Echo +description: + human: + en_US: Return the supplied message. + llm: Return the supplied message. +parameters: + - name: message + type: string + required: true + label: + en_US: Message + human_description: + en_US: Message to return. + llm_description: Message to return. + form: llm + - name: wait_for_peer + type: boolean + required: false + label: + en_US: Wait for another invocation + human_description: + en_US: Synchronize concurrent test invocations. + llm_description: Synchronize concurrent test invocations. + form: llm +extra: + python: + source: tools/echo.py diff --git a/tests/fixtures/plugins/contract_plugin/triggers/__init__.py b/tests/fixtures/plugins/contract_plugin/triggers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/fixtures/plugins/contract_plugin/triggers/constructor.py b/tests/fixtures/plugins/contract_plugin/triggers/constructor.py new file mode 100644 index 00000000..c4f93b14 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/triggers/constructor.py @@ -0,0 +1,99 @@ +from collections.abc import Mapping +from typing import Any +from urllib.parse import urlencode + +from werkzeug import Request + +from dify_plugin.entities import I18nObject, ParameterOption +from dify_plugin.entities.oauth import OAuthCredentials, TriggerOAuthCredentials +from dify_plugin.entities.provider_config import CredentialType +from dify_plugin.entities.trigger import Subscription, UnsubscribeResult +from dify_plugin.interfaces.trigger import TriggerSubscriptionConstructor + + +class ContractConstructor(TriggerSubscriptionConstructor): + def _validate_api_key(self, credentials: Mapping[str, Any]) -> None: + if credentials.get("api_key") != "valid": + msg = "invalid credentials" + raise ValueError(msg) + + def _oauth_get_authorization_url( + self, redirect_uri: str, system_credentials: Mapping[str, Any] + ) -> str: + return "https://provider.invalid/authorize?" + urlencode({ + "redirect_uri": redirect_uri, + "client_id": system_credentials["client_id"], + }) + + def _oauth_get_credentials( + self, redirect_uri: str, system_credentials: Mapping[str, Any], request: Request + ) -> TriggerOAuthCredentials: + del redirect_uri, system_credentials + return TriggerOAuthCredentials( + credentials={ + "access_token": "token-" + request.args["code"], + "refresh_token": "refresh", + } + ) + + def _oauth_refresh_credentials( + self, + redirect_uri: str, + system_credentials: Mapping[str, Any], + credentials: Mapping[str, Any], + ) -> OAuthCredentials: + del redirect_uri, system_credentials + return OAuthCredentials( + credentials={**credentials, "access_token": "refreshed"} + ) + + def _create_subscription( + self, + endpoint: str, + parameters: Mapping[str, Any], + credentials: Mapping[str, Any], + credential_type: CredentialType, + ) -> Subscription: + self.last_credentials = credentials + return Subscription( + endpoint=endpoint, + parameters=parameters, + expires_at=100, + properties={ + "external_id": "subscription", + "credential_type": credential_type.value, + }, + ) + + def _refresh_subscription( + self, + subscription: Subscription, + credentials: Mapping[str, Any], + credential_type: CredentialType, + ) -> Subscription: + self.last_credentials = credentials + self.last_credential_type = credential_type + return subscription.model_copy(update={"expires_at": 200}) + + def _delete_subscription( + self, + subscription: Subscription, + credentials: Mapping[str, Any], + credential_type: CredentialType, + ) -> UnsubscribeResult: + self.last_credentials = credentials + self.last_credential_type = credential_type + return UnsubscribeResult( + success=credentials.get("api_key") != "invalid", + message=subscription.properties["external_id"], + ) + + def _fetch_parameter_options( + self, + parameter: str, + credentials: Mapping[str, Any], + credential_type: CredentialType, + ) -> list[ParameterOption]: + self.last_credentials = credentials + self.last_credential_type = credential_type + return [ParameterOption(value="echo", label=I18nObject(en_us=parameter))] diff --git a/tests/fixtures/plugins/contract_plugin/triggers/dispatch.py b/tests/fixtures/plugins/contract_plugin/triggers/dispatch.py new file mode 100644 index 00000000..36e091cc --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/triggers/dispatch.py @@ -0,0 +1,23 @@ +from triggers.constructor import ContractConstructor +from werkzeug import Request, Response + +from dify_plugin.entities.trigger import EventDispatch, Subscription +from dify_plugin.errors.trigger import TriggerValidationError +from dify_plugin.interfaces.trigger import Trigger + +__all__ = ["ContractConstructor", "ContractTrigger"] + + +class ContractTrigger(Trigger): + def _dispatch_event( + self, subscription: Subscription, request: Request + ) -> EventDispatch: + if subscription.properties.get("reject"): + msg = "invalid webhook" + raise TriggerValidationError(msg) + return EventDispatch( + user_id=request.get_json().get("sender", ""), + events=["echo"], + response=Response("accepted", status=202), + payload=request.get_json(), + ) diff --git a/tests/fixtures/plugins/contract_plugin/triggers/echo.yaml b/tests/fixtures/plugins/contract_plugin/triggers/echo.yaml new file mode 100644 index 00000000..f74f20ef --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/triggers/echo.yaml @@ -0,0 +1,24 @@ +identity: + author: sdk_tests + name: echo + label: + en_US: Echo +description: + en_US: Return the decoded webhook payload. +parameters: + - name: event + type: dynamic-select + label: + en_US: Event +extra: + python: + source: triggers/event.py +output_schema: + type: object + properties: + payload: + type: object + parameters: + type: object + body: + type: string diff --git a/tests/fixtures/plugins/contract_plugin/triggers/event.py b/tests/fixtures/plugins/contract_plugin/triggers/event.py new file mode 100644 index 00000000..6b7244a5 --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/triggers/event.py @@ -0,0 +1,31 @@ +from collections.abc import Mapping +from typing import Any + +from werkzeug import Request + +from dify_plugin.entities import I18nObject, ParameterOption +from dify_plugin.entities.trigger import Variables +from dify_plugin.errors.trigger import EventIgnoreError +from dify_plugin.interfaces.trigger import Event + + +class ContractEvent(Event): + def _on_event( + self, + request: Request, + parameters: Mapping[str, Any], + payload: Mapping[str, Any], + ) -> Variables: + if parameters.get("ignore"): + msg = "event filtered" + raise EventIgnoreError(msg) + return Variables( + variables={ + "payload": payload, + "parameters": parameters, + "body": request.get_data(as_text=True), + } + ) + + def _fetch_parameter_options(self, parameter: str) -> list[ParameterOption]: + return [ParameterOption(value="echo", label=I18nObject(en_us=parameter))] diff --git a/tests/fixtures/plugins/contract_plugin/triggers/events.yaml b/tests/fixtures/plugins/contract_plugin/triggers/events.yaml new file mode 100644 index 00000000..3884940c --- /dev/null +++ b/tests/fixtures/plugins/contract_plugin/triggers/events.yaml @@ -0,0 +1,40 @@ +identity: + author: sdk_tests + name: events + label: + en_US: Events + description: + en_US: Local trigger provider. + icon: icon.svg +subscription_schema: + - name: external_id + type: text-input + label: + en_US: External ID +subscription_constructor: + parameters: + - name: event + type: dynamic-select + label: + en_US: Event + credentials_schema: + - name: api_key + type: secret-input + label: + en_US: API key + oauth_schema: + client_schema: + - name: client_id + type: secret-input + label: + en_US: Client ID + credentials_schema: + - name: access_token + type: secret-input + label: + en_US: Access token +events: + - triggers/echo.yaml +extra: + python: + source: triggers/dispatch.py diff --git a/tests/http_server.py b/tests/http_server.py new file mode 100644 index 00000000..c51f1c88 --- /dev/null +++ b/tests/http_server.py @@ -0,0 +1,48 @@ +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field +from threading import Event + +from pytest_httpserver import HTTPServer as BaseHTTPServer + + +class HTTPServer(BaseHTTPServer): + def thread_target(self) -> None: + # The default poll adds up to 0.5 seconds to every fixture teardown. + assert self.server is not None + self.server.serve_forever(poll_interval=0.01) + + +@dataclass +class Response: + status: int = 200 + headers: dict[str, str] = field(default_factory=dict) + body: bytes = b"" + chunks: tuple[bytes, ...] | None = None + delay: float = 0 + + def stream(self, stopped: Event) -> Iterator[bytes]: + # Send headers before delaying the body, as a real streaming service does. + yield b"" + for chunk in self.chunks if self.chunks is not None else (self.body,): + if stopped.wait(self.delay): + return + yield chunk + + +@dataclass +class Request: + method: str + path: str + headers: Mapping[str, str] + body: bytes + + +@dataclass +class TestHTTPServer: + __test__ = False + + url: str + requests: list[Request] = field(default_factory=list) + + +type HTTPServerFactory = Callable[..., TestHTTPServer] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 1277eb03..7e1a8367 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -1,36 +1,126 @@ -import pathlib -import socket +import json +import os +import shutil +import subprocess # ruff:ignore[suspicious-subprocess-import] import sys +import time +import tomllib from collections.abc import Generator -from typing import ClassVar +from pathlib import Path import pytest -from xprocess import ProcessStarter, XProcess +import yaml +from coverage import Coverage -from ..consts.mockserver import OPENAI_MOCK_SERVER_PORT +from .stdio import StdioPlugin -PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[2] +PROJECT_ROOT = Path(__file__).resolve().parents[2] +PLUGIN_ROOT = PROJECT_ROOT / "tests" / "fixtures" / "plugins" / "contract_plugin" + + +def _package_files(directory: Path) -> set[Path]: + return { + path.relative_to(directory) + for path in directory.rglob("*") + if path.is_file() + and "__pycache__" not in path.parts + and path.suffix not in {".pyc", ".pyo"} + } + + +@pytest.fixture +def local_cli(tmp_path: Path) -> str: + path = tmp_path / "dify-plugin" + source = PROJECT_ROOT / "tests" / "fixtures" / "plugin_cli.py" + path.write_text(f"#!{sys.executable}\n" + source.read_text()) + path.chmod(0o700) + return str(path) @pytest.fixture(scope="session") -def openai_mock_server(xprocess: XProcess) -> Generator[str, None, None]: - class Starter(ProcessStarter): - pattern = "OpenAI mock server starting" - args: ClassVar[list[str]] = [sys.executable, "-m", "tests.__mock_server"] - max_read_lines = 20 - timeout = 30 - terminate_on_interrupt = True - popen_kwargs: ClassVar[dict[str, str]] = {"cwd": str(PROJECT_ROOT)} - - def startup_check(self) -> bool: - try: - with socket.create_connection( - ("127.0.0.1", OPENAI_MOCK_SERVER_PORT), timeout=1 - ): - return True - except OSError: - return False - - xprocess.ensure("openai-mock-server", Starter) - yield f"http://127.0.0.1:{OPENAI_MOCK_SERVER_PORT}" - xprocess.getinfo("openai-mock-server").terminate() +def installed_sdk(tmp_path_factory: pytest.TempPathFactory) -> Path: + build_dir = tmp_path_factory.mktemp("sdk-wheel") + uv = shutil.which("uv") + assert uv is not None, "uv is required to build and install the SDK under test" + subprocess.run( # ruff:ignore[subprocess-without-shell-equals-true] + [uv, "build", "--wheel", "--offline", "--out-dir", str(build_dir)], + cwd=PROJECT_ROOT, + check=True, + capture_output=True, + text=True, + timeout=60, + ) + wheels = list(build_dir.glob("*.whl")) + assert len(wheels) == 1 + install_dir = build_dir / "site-packages" + subprocess.run( # ruff:ignore[subprocess-without-shell-equals-true] + [ + uv, + "pip", + "install", + "--python", + sys.executable, + "--offline", + "--no-deps", + "--target", + str(install_dir), + str(wheels[0]), + ], + check=True, + capture_output=True, + text=True, + timeout=60, + ) + source_package = PROJECT_ROOT / "src" / "dify_plugin" + installed_package = install_dir / "dify_plugin" + source_files = _package_files(source_package) + assert _package_files(installed_package) == source_files + for relative_path in source_files: + assert (installed_package / relative_path).read_bytes() == ( + source_package / relative_path + ).read_bytes(), relative_path + return install_dir + + +@pytest.fixture +def stdio_plugin(installed_sdk: Path, tmp_path: Path) -> Generator[StdioPlugin]: + plugin_dir = tmp_path / "plugin" + shutil.copytree( + PLUGIN_ROOT, plugin_dir, ignore=shutil.ignore_patterns("__pycache__") + ) + manifest_path = plugin_dir / "manifest.yaml" + manifest = yaml.safe_load(manifest_path.read_text()) + # The daemon permits tools and endpoints together, but other kinds are exclusive. + manifest["plugins"] = { + kind: manifest["plugins"][kind] for kind in ("tools", "endpoints") + } + manifest_path.write_text(yaml.safe_dump(manifest)) + stderr_path = tmp_path / "plugin.stderr" + env = {**os.environ, "PYTHONPATH": str(installed_sdk), "PYTHONNOUSERSITE": "1"} + env.pop("PYTHONUNBUFFERED", None) + with stderr_path.open("wb") as stderr: + process = subprocess.Popen( + [sys.executable, "-m", "main"], + cwd=plugin_dir, + env=env, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=stderr, + ) + plugin = StdioPlugin(process, stderr_path) + try: + deadline = time.monotonic() + 10 + plugin.manifest = plugin.read_message(deadline) + assert plugin.manifest["name"] == "contract_plugin" + while plugin.read_message(deadline).get("event") != "heartbeat": + pass + identity = json.loads(stderr_path.read_text().splitlines()[0]) + sdk_file = Path(identity["sdk_file"]) + assert sdk_file == installed_sdk / "dify_plugin" / "__init__.py" + project = tomllib.loads((PROJECT_ROOT / "pyproject.toml").read_text()) + assert identity["sdk_version"] == project["project"]["version"] + if Coverage.current() is not None: + assert identity["coverage_active"], "Plugin coverage did not start" + yield plugin + finally: + plugin.close() diff --git a/tests/integration/stdio.py b/tests/integration/stdio.py new file mode 100644 index 00000000..f3df5004 --- /dev/null +++ b/tests/integration/stdio.py @@ -0,0 +1,98 @@ +import json +import os +import select +import subprocess # ruff:ignore[suspicious-subprocess-import] +import time +from pathlib import Path +from typing import Any + + +class StdioPlugin: + def __init__(self, process: subprocess.Popen[bytes], stderr_path: Path) -> None: + self.process = process + self.stderr_path = stderr_path + self.stdout_path = stderr_path.with_suffix(".stdout") + self.stdout_path.touch() + self.manifest: dict[str, Any] = {} + self.pending = b"" + + def read_message(self, deadline: float | None = None) -> dict[str, Any]: + assert self.process.stdout is not None + if deadline is None: + deadline = time.monotonic() + 10 + while True: + if b"\n" in self.pending: + line, self.pending = self.pending.split(b"\n", 1) + if line.strip(): + return json.loads(line) + continue + remaining = deadline - time.monotonic() + if ( + remaining <= 0 + or not select.select([self.process.stdout], [], [], remaining)[0] + ): + msg = ( + f"Plugin response timed out. stdout: {self.stdout_path}; stderr:\n" + + self.stderr_path.read_text() + ) + raise AssertionError(msg) + data = os.read(self.process.stdout.fileno(), 65536) + if not data: + msg = f"Plugin stdout closed. stderr:\n{self.stderr_path.read_text()}" + raise AssertionError(msg) + with self.stdout_path.open("ab") as output: + output.write(data) + self.pending += data + + def send(self, *messages: dict[str, Any]) -> None: + self.send_bytes( + b"".join( + json.dumps(message, ensure_ascii=False).encode() + b"\n" + for message in messages + ) + ) + + def send_bytes(self, data: bytes) -> None: + assert self.process.stdin is not None + deadline = time.monotonic() + 10 + offset = 0 + while offset < len(data): + remaining = deadline - time.monotonic() + if ( + remaining <= 0 + or not select.select([], [self.process.stdin], [], remaining)[1] + ): + msg = "Plugin stopped reading stdin" + raise AssertionError(msg) + offset += os.write( + self.process.stdin.fileno(), data[offset : offset + 4096] + ) + + def receive_sessions(self, *session_ids: str) -> dict[str, list[dict[str, Any]]]: + results: dict[str, list[dict[str, Any]]] = { + session_id: [] for session_id in session_ids + } + finished: set[str] = set() + deadline = time.monotonic() + 10 + while len(finished) < len(session_ids): + message = self.read_message(deadline) + if message.get("event") != "session": + continue + session_id = message["session_id"] + assert session_id in results, message + assert session_id not in finished, message + results[session_id].append(message["data"]) + if message["data"]["type"] == "end": + finished.add(session_id) + return results + + def close(self) -> None: + if self.process.poll() is None: + self.process.terminate() + try: + self.process.communicate(timeout=5) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.communicate(timeout=5) + msg = "Plugin did not exit after SIGTERM" + raise AssertionError(msg) from None diff --git a/tests/integration/test_invoke_llm.py b/tests/integration/test_invoke_llm.py deleted file mode 100644 index 7312c203..00000000 --- a/tests/integration/test_invoke_llm.py +++ /dev/null @@ -1,65 +0,0 @@ -import pathlib - -import pytest - -from dify_plugin.config.integration_config import IntegrationConfig, find_dify_cli_path -from dify_plugin.core.entities.plugin.request import ( - ModelActions, - ModelInvokeLLMRequest, - PluginInvokeType, -) -from dify_plugin.entities.model import ModelType -from dify_plugin.entities.model.llm import LLMResultChunk -from dify_plugin.entities.model.message import UserPromptMessage -from dify_plugin.integration.run import PluginRunner - -# Import requests only after dify_plugin applies its gevent patch. -# isort: split - -import requests - -_OPENAI_PLUGIN_URL = ( - "https://marketplace.dify.ai/api/v1/plugins/langgenius/openai/1.0.0/download" -) - -pytestmark = pytest.mark.skipif( - find_dify_cli_path() is None, - reason="dify cli not found; install dify-plugin-cli to run integration tests", -) - - -def test_invoke_llm(openai_mock_server: str) -> None: - response = requests.get(_OPENAI_PLUGIN_URL, timeout=10) - response.raise_for_status() - - # save the response to a file - pathlib.Path("langgenius-openai.difypkg").write_bytes(response.content) - - # run the plugin - with PluginRunner( - config=IntegrationConfig(), plugin_package_path="langgenius-openai.difypkg" - ) as runner: - for result in runner.invoke( - access_type=PluginInvokeType.Model, - access_action=ModelActions.InvokeLLM, - payload=ModelInvokeLLMRequest( - prompt_messages=[ - UserPromptMessage(content="Hello, world!"), - ], - user_id="", - provider="openai", - model_type=ModelType.LLM, - model="gpt-5.6", - credentials={ - "api_protocol": "chat", - "openai_api_base": openai_mock_server, - "openai_api_key": "test", - }, - model_parameters={}, - stop=[], - tools=[], - stream=False, - ), - response_type=LLMResultChunk, - ): - assert result.delta.message.content == "Hello, world!" diff --git a/tests/integration/test_plugin_runner.py b/tests/integration/test_plugin_runner.py new file mode 100644 index 00000000..9b76736f --- /dev/null +++ b/tests/integration/test_plugin_runner.py @@ -0,0 +1,411 @@ +import os +import threading +from collections.abc import Generator +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import gevent +import pytest +from pydantic import BaseModel, RootModel + +from dify_plugin.config.integration_config import IntegrationConfig +from dify_plugin.core.entities.plugin.request import ( + ModelActions, + ModelInvokeRerankRequest, + PluginAccessAction, + PluginInvokeType, + ToolActions, + ToolInvokeRequest, +) +from dify_plugin.integration import run as runner_module +from dify_plugin.integration.exc import PluginStoppedError +from dify_plugin.integration.run import PluginRunner + +pytestmark = pytest.mark.integration + + +class Payload(BaseModel): + value: str + + +def _invoke(runner: PluginRunner, value: str = "hello") -> Generator[Payload]: + return runner.invoke( + PluginInvokeType.Tool, + ToolActions.InvokeTool, + payload=Payload(value=value), + response_type=Payload, + ) + + +def _assert_pipes_closed(runner: PluginRunner) -> None: + # gevent defers os.close until the next event loop iteration. + gevent.sleep(0.001) + for pipe in ( + runner.stdout_pipe_read, + runner.stdout_pipe_write, + runner.stderr_pipe_read, + runner.stderr_pipe_write, + runner.stdin_pipe_read, + runner.stdin_pipe_write, + ): + with pytest.raises(OSError, match="Bad file descriptor"): + os.fstat(pipe) + + +def test_child_exit_unblocks_pending_invocation(local_cli: str) -> None: + with ( + PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", "exit-on-invoke"], + ) as runner, + gevent.Timeout(2, AssertionError("Invocation did not notice child exit")), + pytest.raises(PluginStoppedError), + ): + list(_invoke(runner)) + + +@pytest.mark.parametrize("scenario", ["echo", "mixed", "stderr", "stderr-error"]) +def test_runner_reads_local_cli_protocol(local_cli: str, scenario: str) -> None: + read_pipes: set[int] = set() + + class Runner(PluginRunner): + def _message_reader(self, pipe: int) -> None: + assert pipe == self.stdout_pipe_read + super()._message_reader(pipe) + + def _read_async(self, fd: int) -> bytes: + read_pipes.add(fd) + return super()._read_async(fd) + + with ( + Runner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", scenario], + ) as runner, + gevent.Timeout(2, AssertionError("Invocation stopped responding")), + ): + assert list(_invoke(runner, "你好")) == [Payload(value="你好")] + assert read_pipes == {runner.stdout_pipe_read} + assert runner.process.stdin is None + assert runner.process.stdout is None + assert runner.process.stderr is None + assert runner.process.poll() is not None + assert not runner.stdout_reader.is_alive() + _assert_pipes_closed(runner) + + +@pytest.mark.parametrize( + ("access_type", "access_action", "payload_type", "payload"), + [ + ( + PluginInvokeType.Tool, + ToolActions.InvokeTool, + ToolInvokeRequest, + { + "type": "tool", + "action": "invoke_tool", + "user_id": "tool-user", + "provider": "json_process", + "tool": "parse", + "credentials": {"api_key": "tool-key"}, + "credential_type": "api-key", + "tool_parameters": { + "content": '{"message":"你好"}', + "ensure_ascii": False, + "limit": 0, + }, + }, + ), + ( + PluginInvokeType.Model, + ModelActions.InvokeRerank, + ModelInvokeRerankRequest, + { + "type": "model", + "action": "invoke_rerank", + "user_id": "model-user", + "provider": "openrouter", + "model_type": "rerank", + "model": "test-reranker", + "credentials": {"api_key": "model-key"}, + "query": "capital of France", + "docs": ["Berlin", "Paris"], + "score_threshold": None, + "top_n": 1, + }, + ), + ], + ids=["tool", "model"], +) +def test_runner_preserves_cli_request_fields( + local_cli: str, + access_type: PluginInvokeType, + access_action: PluginAccessAction, + payload_type: type[BaseModel], + payload: dict, +) -> None: + with ( + PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", "record-request"], + ) as runner, + gevent.Timeout(2, AssertionError("Invocation stopped responding")), + ): + (response,) = runner.invoke( + access_type, + access_action, + payload=payload_type.model_validate(payload), + response_type=RootModel[dict], + ) + wire = response.root + assert wire.pop("invoke_id") + assert wire == { + "type": payload["type"], + "action": payload["action"], + "request": payload, + } + + +def test_runner_routes_interleaved_invocations(local_cli: str) -> None: + with ( + ThreadPoolExecutor(max_workers=2) as executor, + PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", "reverse"], + ) as runner, + gevent.Timeout(2, AssertionError("Concurrent invocations stopped responding")), + ): + results = list( + executor.map(lambda text: list(_invoke(runner, text)), ["a", "b"]) + ) + assert results == [[Payload(value="a")], [Payload(value="b")]] + + +@pytest.mark.timeout(5) +def test_timeout_closes_runner_before_waiting_for_invocation_workers( + local_cli: str, +) -> None: + with ( + pytest.raises(TimeoutError, match="No invocation response"), + ThreadPoolExecutor(max_workers=2) as executor, + PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", "no-response"], + ) as runner, + gevent.Timeout(0.1, TimeoutError("No invocation response")), + ): + list(executor.map(lambda value: list(_invoke(runner, value)), ["a", "b"])) + + assert runner.process.poll() is not None + assert not runner.stdout_reader.is_alive() + assert runner.q == {} + + +def test_runner_handles_large_requests_and_closed_streams(local_cli: str) -> None: + with PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), "fixture.difypkg" + ) as runner: + text = "消息" * 20_000 + with gevent.Timeout(2, AssertionError("Large invocation stopped responding")): + stream = _invoke(runner, text) + assert next(stream) == Payload(value=text) + stream.close() + assert runner.q == {} + assert list(_invoke(runner)) == [Payload(value="hello")] + with pytest.raises(PluginStoppedError): + list(_invoke(runner)) + + +@pytest.mark.parametrize( + ("scenario", "message"), + [ + ("error-on-invoke", "invocation failed"), + ("invalid-response", "Invalid response type"), + ], +) +def test_runner_propagates_cli_errors( + local_cli: str, scenario: str, message: str +) -> None: + with PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", scenario], + ) as runner: + with ( + gevent.Timeout(2, AssertionError("Invocation stopped responding")), + pytest.raises(ValueError, match=message), + ): + list(_invoke(runner)) + assert runner.q == {} + + +def test_runner_cleans_packaged_source_directory( + local_cli: str, tmp_path: Path +) -> None: + source = tmp_path / "source" + source.mkdir() + with PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), str(source) + ) as runner: + package = Path(runner.plugin_package_path) + assert package.read_bytes() == b"local fixture package" + assert list(_invoke(runner)) == [Payload(value="hello")] + assert not package.parent.exists() + + +@pytest.mark.parametrize( + "scenario", ["exit-before-ready", "error-before-ready", "never-ready"] +) +def test_startup_failure_preserves_timeout_and_reaps_process( + local_cli: str, + scenario: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + processes = [] + pipes = [] + pipe = runner_module.os.pipe + popen = runner_module.subprocess.Popen + + def record_pipe() -> tuple[int, int]: + descriptors = pipe() + pipes.extend(descriptors) + return descriptors + + def record_process( + *args: object, **kwargs: object + ) -> runner_module.subprocess.Popen[bytes]: + process = popen(*args, **kwargs) + processes.append(process) + return process + + config = IntegrationConfig(dify_cli_path=local_cli, _env_file=None) + monkeypatch.setattr(runner_module.os, "pipe", record_pipe) + monkeypatch.setattr(runner_module.subprocess, "Popen", record_process) + semaphore = runner_module.Semaphore(0) + acquire = semaphore.acquire + + def acquire_ready(*, timeout: float) -> bool: + assert timeout == 30 + return acquire(timeout=0.2) + + monkeypatch.setattr(semaphore, "acquire", acquire_ready) + monkeypatch.setattr(runner_module, "Semaphore", lambda _value: semaphore) + with ( + gevent.Timeout(2, AssertionError("Constructor did not notice child exit")), + pytest.raises(TimeoutError, match="Plugin failed to start within 30 seconds"), + ): + PluginRunner(config, "fixture.difypkg", extra_args=["--scenario", scenario]) + assert len(processes) == 1 + assert processes[0].poll() is not None + gevent.sleep(0.001) + for fd in pipes[:6]: + with pytest.raises(OSError, match="Bad file descriptor"): + os.fstat(fd) + + +def test_packaging_failure_removes_temporary_directory( + local_cli: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + temporary_paths = [] + temporary_directory = runner_module.tempfile.TemporaryDirectory + + def record_directory() -> runner_module.tempfile.TemporaryDirectory[str]: + directory = temporary_directory() + temporary_paths.append(Path(directory.name)) + return directory + + source = tmp_path / "fail-package" + source.mkdir() + monkeypatch.setattr(runner_module.tempfile, "TemporaryDirectory", record_directory) + with pytest.raises(runner_module.subprocess.CalledProcessError): + PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), str(source) + ) + assert len(temporary_paths) == 1 + assert not temporary_paths[0].exists() + + +def test_runner_waits_for_graceful_shutdown( + local_cli: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + gate = tmp_path / "allow-shutdown" + runner = PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), + "fixture.difypkg", + extra_args=["--scenario", "graceful-termination", str(gate)], + ) + wait = runner.process.wait + + def wait_for_exit(timeout: float | None = None) -> int: + assert timeout is None + return wait(timeout=timeout) + + monkeypatch.setattr(runner.process, "wait", wait_for_exit) + closing = gevent.spawn(runner._close) + try: + with gevent.Timeout(2): + while not gate.with_suffix(".started").exists(): + gevent.sleep(0.01) + assert not closing.ready() + assert runner.process.poll() is None + finally: + gate.touch() + closing.get(timeout=2) + assert gate.with_suffix(".completed").exists() + assert runner.process.poll() is not None + assert not runner.stdout_reader.is_alive() + + +def test_concurrent_context_exits_wait_for_resource_cleanup( + local_cli: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + source = tmp_path / "source" + source.mkdir() + runner = PluginRunner( + IntegrationConfig(dify_cli_path=local_cli, _env_file=None), str(source) + ) + package = Path(runner.plugin_package_path) + cleanup_started = threading.Event() + allow_cleanup = threading.Event() + wait = runner.process.wait + + def delayed_wait(**kwargs: object) -> int: + cleanup_started.set() + assert allow_cleanup.wait(timeout=2) + return wait(**kwargs) + + monkeypatch.setattr(runner.process, "wait", delayed_wait) + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(runner.__exit__, None, None, None) + assert cleanup_started.wait(timeout=2) + second = executor.submit(runner.__exit__, None, None, None) + try: + with pytest.raises(TimeoutError): + second.result(timeout=0.05) + finally: + allow_cleanup.set() + first.result(timeout=2) + second.result(timeout=2) + assert runner.process.poll() is not None + assert not runner.stdout_reader.is_alive() + _assert_pipes_closed(runner) + assert not package.parent.exists() + + +def test_reader_start_failure_keeps_original_error( + local_cli: str, monkeypatch: pytest.MonkeyPatch +) -> None: + config = IntegrationConfig(dify_cli_path=local_cli, _env_file=None) + + def fail_start(_thread: threading.Thread) -> None: + msg = "reader start failed" + raise RuntimeError(msg) + + monkeypatch.setattr(runner_module.threading.Thread, "start", fail_start) + with pytest.raises(RuntimeError, match="reader start failed"): + PluginRunner(config, "fixture.difypkg") diff --git a/tests/integration/test_stdio_plugin.py b/tests/integration/test_stdio_plugin.py new file mode 100644 index 00000000..bd2a1028 --- /dev/null +++ b/tests/integration/test_stdio_plugin.py @@ -0,0 +1,92 @@ +from typing import Any + +import pytest + +from .stdio import StdioPlugin + +pytestmark = pytest.mark.integration + + +def _invoke(session_id: str, message: str) -> dict[str, Any]: + return { + "session_id": session_id, + "event": "request", + "data": { + "type": "tool", + "action": "invoke_tool", + "user_id": "test-user", + "provider": "contract", + "tool": "echo", + "credentials": {"api_key": "valid"}, + "tool_parameters": {"message": message}, + }, + } + + +def _assert_echo(events: list[dict[str, Any]], message: str) -> None: + assert [event["type"] for event in events] == ["stream", "stream", "end"] + assert events[0]["data"]["type"] == "text" + assert events[0]["data"]["message"] == {"text": message} + assert events[1]["data"]["type"] == "json" + assert events[1]["data"]["message"] == {"json_object": {"echo": message}} + assert events[2]["data"] == {} + + +@pytest.mark.parametrize( + "message", ["Hello, 世界 🌍", "消息" * 20_000], ids=["unicode", "large-message"] +) +def test_installed_wheel_registers_and_invokes_over_stdio( + stdio_plugin: StdioPlugin, message: str +) -> None: + assert "provider/contract.yaml" in stdio_plugin.manifest["plugins"]["tools"] + + stdio_plugin.send(_invoke("echo", message)) + + _assert_echo(stdio_plugin.receive_sessions("echo")["echo"], message) + + +def test_stdio_keeps_concurrent_sessions_separate(stdio_plugin: StdioPlugin) -> None: + requests = [_invoke("first", "one"), _invoke("second", "two")] + for request in requests: + request["data"]["tool_parameters"]["wait_for_peer"] = True + stdio_plugin.send(*requests) + + responses = stdio_plugin.receive_sessions("first", "second") + + _assert_echo(responses["first"], "one") + _assert_echo(responses["second"], "two") + + +def test_stdio_reports_call_error_and_accepts_next_request( + stdio_plugin: StdioPlugin, +) -> None: + request = _invoke("missing-tool", "unused") + request["data"]["tool"] = "missing" + stdio_plugin.send(request) + + events = stdio_plugin.receive_sessions("missing-tool")["missing-tool"] + + assert [event["type"] for event in events] == ["error", "end"] + assert events[0]["data"]["error_type"] == "ValueError" + assert "Tool `missing` not found" in events[0]["data"]["message"] + stdio_plugin.send(_invoke("recovered", "ok")) + _assert_echo(stdio_plugin.receive_sessions("recovered")["recovered"], "ok") + + +def test_stdio_recovers_from_malformed_json(stdio_plugin: StdioPlugin) -> None: + stdio_plugin.send_bytes(b"{invalid json}\n") + + error = stdio_plugin.read_message() + + assert error["event"] == "error" + assert error["data"]["error"] + stdio_plugin.send(_invoke("recovered", "ok")) + _assert_echo(stdio_plugin.receive_sessions("recovered")["recovered"], "ok") + + +def test_stdio_process_exits_on_termination(stdio_plugin: StdioPlugin) -> None: + assert stdio_plugin.process.poll() is None + + stdio_plugin.close() + + assert stdio_plugin.process.poll() is not None diff --git a/tests/interfaces/agent/test_agent_behavior.py b/tests/interfaces/agent/test_agent_behavior.py new file mode 100644 index 00000000..d0f19f8f --- /dev/null +++ b/tests/interfaces/agent/test_agent_behavior.py @@ -0,0 +1,375 @@ +import signal +from collections.abc import Generator +from copy import deepcopy +from decimal import Decimal + +import pytest +from pydantic import ValidationError + +from dify_plugin.core.runtime import Session +from dify_plugin.entities import I18nObject +from dify_plugin.entities.agent import AgentInvokeMessage, AgentRuntime +from dify_plugin.entities.model import ( + AIModelEntity, + ModelPropertyKey, + ModelType, + ParameterRule, +) +from dify_plugin.entities.model.llm import LLMUsage +from dify_plugin.entities.model.message import ( + AssistantPromptMessage, + PromptMessage, + PromptMessageTool, + TextPromptMessageContent, + ToolPromptMessage, + UserPromptMessage, +) +from dify_plugin.entities.tool import ToolParameter +from dify_plugin.file.constants import DIFY_FILE_IDENTITY +from dify_plugin.interfaces.agent import ( + AgentModelConfig, + AgentStrategy, + AgentToolIdentity, + ToolEntity, +) + + +class SchemaAgent(AgentStrategy): + def _invoke(self, parameters: dict) -> Generator[AgentInvokeMessage]: + for file in parameters.get("files", []): + yield self.create_variable_message("filename", file.filename) + for tool in self._init_prompt_tools(parameters.get("tools")): + yield self.create_json_message(tool.model_dump()) + + +@pytest.fixture +def agent() -> Generator[SchemaAgent]: + original_sigint = signal.getsignal(signal.SIGINT) + session = Session.empty_session() + try: + yield SchemaAgent(runtime=AgentRuntime(user_id="user"), session=session) + finally: + session.reader.close() + session.writer.close() + session._executor.shutdown(wait=True) + signal.signal(signal.SIGINT, original_sigint) + + +def make_tool(name: str, parameters: list[ToolParameter] | None = None) -> ToolEntity: + return ToolEntity( + identity=AgentToolIdentity( + author="sdk_tests", + name=name, + label=I18nObject(en_us=name), + provider="contract", + ), + parameters=parameters or [], + ) + + +def test_agent_public_invoke_converts_files_and_builds_only_llm_parameters( + agent: SchemaAgent, +) -> None: + label = I18nObject(en_us="Input") + parameters = [ + ToolParameter( + name=name, + type=kind, + form=form, + required=True, + label=label, + human_description=label, + llm_description="Input value", + ) + for name, kind, form in [ + ("count", "number", "llm"), + ("enabled", "boolean", "llm"), + ("settings", "object", "llm"), + ("items", "array", "llm"), + ("key", "secret-input", "form"), + ("schema", "string", "schema"), + ("file", "file", "llm"), + ("files", "files", "llm"), + ] + ] + messages = list( + agent.invoke({ + "files": [ + { + "dify_model_identity": DIFY_FILE_IDENTITY, + "url": "https://files.invalid/a.txt", + "type": "document", + "filename": "a.txt", + } + ], + "tools": [make_tool("configure", parameters)], + }) + ) + + assert all(isinstance(message, AgentInvokeMessage) for message in messages) + assert messages[0].message.variable_value == "a.txt" + schema = messages[1].message.json_object + assert schema["name"] == "configure" + assert not schema["description"] + assert schema["parameters"] == { + "type": "object", + "properties": { + "count": {"type": "number", "description": "Input value"}, + "enabled": {"type": "boolean", "description": "Input value"}, + "settings": {"type": "object", "description": "Input value"}, + "items": {"type": "array", "description": "Input value"}, + }, + "required": ["count", "enabled", "settings", "items"], + } + + +def test_one_invalid_tool_is_diagnosed_without_dropping_valid_tools( + agent: SchemaAgent, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + convert = agent._convert_tool_to_prompt_message_tool + + def convert_with_removed_tool(tool: ToolEntity) -> PromptMessageTool: + if tool.identity.name == "removed": + msg = "provider was removed" + raise ValueError(msg) + return convert(tool) + + monkeypatch.setattr( + agent, "_convert_tool_to_prompt_message_tool", convert_with_removed_tool + ) + messages = list( + agent.invoke({ + "tools": [make_tool("before"), make_tool("removed"), make_tool("after")] + }) + ) + assert [message.message.json_object["name"] for message in messages] == [ + "before", + "after", + ] + assert len(caplog.records) == 1 + assert "removed" in caplog.records[0].message + assert "conversion failed" in caplog.records[0].message + assert isinstance(caplog.records[0].exc_info[1], ValueError) + assert list(agent.invoke({"tools": None})) == [] + + +def test_usage_accumulation_preserves_existing_shared_result_behavior( + agent: SchemaAgent, +) -> None: + first = LLMUsage.empty_usage().model_copy( + update={ + "prompt_tokens": 10, + "completion_tokens": 4, + "total_tokens": 14, + "prompt_price": Decimal("0.01"), + "completion_price": Decimal("0.008"), + "total_price": Decimal("0.018"), + } + ) + second = LLMUsage.empty_usage().model_copy( + update={ + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + "prompt_price": Decimal("0.003"), + "completion_price": Decimal("0.004"), + "total_price": Decimal("0.007"), + } + ) + total: dict[str, LLMUsage | None] = {"usage": None} + agent.increase_usage(total, first) + agent.increase_usage(total, second) + assert total["usage"] is first + assert total["usage"].prompt_tokens == 13 + assert total["usage"].completion_tokens == 6 + assert total["usage"].total_tokens == 19 + assert total["usage"].prompt_price == Decimal("0.013") + assert total["usage"].completion_price == Decimal("0.012") + assert total["usage"].total_price == Decimal("0.025") + + agent.increase_usage(total, first) + assert total["usage"].prompt_tokens == 26 + assert total["usage"].completion_tokens == 12 + assert total["usage"].total_tokens == 38 + assert total["usage"].prompt_price == Decimal("0.026") + assert total["usage"].completion_price == Decimal("0.024") + assert total["usage"].total_price == Decimal("0.050") + + +@pytest.mark.parametrize("parameter_name", ["max_tokens", "max_completion_tokens"]) +@pytest.mark.parametrize(("requested", "expected"), [(30, 30), (60, 60), (80, 60)]) +def test_max_tokens_respects_context_limit_and_provider_parameter_name( + agent: SchemaAgent, + monkeypatch: pytest.MonkeyPatch, + parameter_name: str, + requested: int, + expected: int, +) -> None: + model = AIModelEntity( + model="test", + model_type=ModelType.LLM, + model_properties={ModelPropertyKey.CONTEXT_SIZE: 100}, + parameter_rules=[ParameterRule(name=parameter_name, use_template="max_tokens")], + ) + prompts = [UserPromptMessage(content="model input")] + + def count_tokens(messages: list[PromptMessage]) -> int: + assert messages == prompts + return 40 + + monkeypatch.setattr(agent, "_get_num_tokens_by_gpt2", count_tokens) + parameters = {parameter_name: requested, "temperature": 0.2} + agent.recalc_llm_max_tokens(model, prompts, parameters) + assert parameters == {parameter_name: expected, "temperature": 0.2} + + +def test_model_without_context_limit_keeps_request_and_skips_tokenization( + agent: SchemaAgent, monkeypatch: pytest.MonkeyPatch +) -> None: + model = AIModelEntity( + model="unbounded", model_type=ModelType.LLM, model_properties={} + ) + + def reject_tokenization(_messages: list[PromptMessage]) -> int: + pytest.fail("a model without a context limit must not tokenize the prompt") + + monkeypatch.setattr(agent, "_get_num_tokens_by_gpt2", reject_tokenization) + parameters = {"max_tokens": 100} + agent.recalc_llm_max_tokens(model, [], parameters) + assert parameters == {"max_tokens": 100} + + +@pytest.mark.parametrize("remaining", [15, 1, 0, -1]) +def test_context_estimate_keeps_the_provider_call_possible( + agent: SchemaAgent, monkeypatch: pytest.MonkeyPatch, remaining: int +) -> None: + model = AIModelEntity( + model="test", + model_type=ModelType.LLM, + model_properties={ModelPropertyKey.CONTEXT_SIZE: 100}, + parameter_rules=[ParameterRule(name="max_tokens", use_template="max_tokens")], + ) + monkeypatch.setattr(agent, "_get_num_tokens_by_gpt2", lambda _: 100 - remaining) + parameters = {"max_tokens": 16} + + agent.recalc_llm_max_tokens(model, [], parameters) + + assert parameters == {"max_tokens": 16} + + +def test_model_history_reconstructs_tool_calls_and_multimodal_content() -> None: + # The official ReAct and FunctionCalling strategies parse this Dify payload. + data = { + "provider": "openai", + "model": "gpt-4o", + "mode": "chat", + "history_prompt_messages": [ + { + "role": "user", + "content": [{"type": "text", "data": "Search documentation"}], + }, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "search-1", + "type": "function", + "function": {"name": "search", "arguments": '{"query":"SDK"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "search-1", "content": "Result"}, + ], + } + original = deepcopy(data) + + first = AgentModelConfig.model_validate(data) + second = AgentModelConfig.model_validate(data) + user, assistant, tool = first.history_prompt_messages + + assert isinstance(user, UserPromptMessage) + assert user.get_text_content() == "Search documentation" + assert isinstance(assistant, AssistantPromptMessage) + assert isinstance(tool, ToolPromptMessage) + assert assistant.tool_calls[0].id == tool.tool_call_id == "search-1" + assert assistant.tool_calls[0].function.arguments == '{"query":"SDK"}' + assert first == second + assert data == original + + +@pytest.mark.parametrize( + "history", + [ + {"role": "user", "content": "hello"}, + [{"role": "tool", "content": "missing call ID"}], + ], +) +def test_model_history_rejects_malformed_history(history: object) -> None: + with pytest.raises((TypeError, ValidationError)): + AgentModelConfig.model_validate({ + "provider": "openai", + "model": "gpt-4o", + "mode": "chat", + "history_prompt_messages": history, + }) + + +@pytest.mark.parametrize("multipart", [False, True]) +def test_real_tokenizer_preserves_budget_for_each_content_form( + agent: SchemaAgent, multipart: bool +) -> None: + model = AIModelEntity( + model="test", + model_type=ModelType.LLM, + model_properties={ModelPropertyKey.CONTEXT_SIZE: 100}, + parameter_rules=[ParameterRule(name="max_tokens", use_template="max_tokens")], + ) + text = "hello " * 50 + content = [TextPromptMessageContent(data=text)] if multipart else text + parameters = {"max_tokens": 80} + + agent.recalc_llm_max_tokens(model, [UserPromptMessage(content=content)], parameters) + + assert parameters["max_tokens"] == (80 if multipart else 49) + + +def test_gpt2_overestimate_does_not_reject_a_valid_provider_prompt( + agent: SchemaAgent, +) -> None: + # GPT-4o encodes this text in 48,000 tokens; GPT-2 estimates 128,000. + text = "今天天气很好\uff0c我们一起去公园散步。" * 4000 + model = AIModelEntity( + model="gpt-4o", + model_type=ModelType.LLM, + model_properties={ModelPropertyKey.CONTEXT_SIZE: 128000}, + parameter_rules=[ParameterRule(name="max_tokens", use_template="max_tokens")], + ) + parameters = {"max_tokens": 1024} + + agent.recalc_llm_max_tokens(model, [UserPromptMessage(content=text)], parameters) + + assert parameters["max_tokens"] > 0 + + +def test_real_tokenizer_preserves_special_token_error( + agent: SchemaAgent, +) -> None: + model = AIModelEntity( + model="test", + model_type=ModelType.LLM, + model_properties={ModelPropertyKey.CONTEXT_SIZE: 100}, + parameter_rules=[ParameterRule(name="max_tokens", use_template="max_tokens")], + ) + parameters = {"max_tokens": 1} + + with pytest.raises(ValueError, match="special token"): + agent.recalc_llm_max_tokens( + model, + [UserPromptMessage(content="Explain <|endoftext|> as ordinary text")], + parameters, + ) + + assert parameters == {"max_tokens": 1} diff --git a/tests/interfaces/datasource/test_construct_datasource.py b/tests/interfaces/datasource/test_construct_datasource.py deleted file mode 100644 index b184ab80..00000000 --- a/tests/interfaces/datasource/test_construct_datasource.py +++ /dev/null @@ -1,85 +0,0 @@ -from collections.abc import Generator, Mapping -from concurrent.futures import ThreadPoolExecutor -from typing import Any - -from dify_plugin.core.runtime import Session -from dify_plugin.core.server.stdio.request_reader import StdioRequestReader -from dify_plugin.core.server.stdio.response_writer import StdioResponseWriter -from dify_plugin.entities.datasource import ( - DatasourceGetPagesResponse, - DatasourceMessage, - DatasourceRuntime, - GetOnlineDocumentPageContentRequest, -) -from dify_plugin.interfaces.datasource.online_document import OnlineDocumentDatasource -from dify_plugin.interfaces.datasource.website import WebsiteCrawlDatasource - - -def test_construct_website_crawl_datasource() -> None: - """ - Test WebsiteCrawlDatasource can be constructed in specific session - """ - - class Website(WebsiteCrawlDatasource): - def _get_website_crawl( - self, datasource_parameters: Mapping[str, Any] - ) -> Generator[DatasourceMessage, None, None]: - yield DatasourceMessage( - type=DatasourceMessage.MessageType.TEXT, - message=DatasourceMessage.TextMessage( - text=( - "Website crawl result for " - f"{datasource_parameters.get('url', 'unknown')}" - ) - ), - ) - - session = Session( - session_id="test", - executor=ThreadPoolExecutor(max_workers=1), - reader=StdioRequestReader(), - writer=StdioResponseWriter(), - ) - datasource = Website( - runtime=DatasourceRuntime(credentials={}, user_id="test", session_id="test"), - session=session, - ) - assert datasource is not None - - -def test_construct_online_document_datasource() -> None: - """ - Test OnlineDocumentDatasource can be constructed in specific session - """ - - class OnlineDocument(OnlineDocumentDatasource): - def _get_pages( - self, datasource_parameters: Mapping[str, Any] - ) -> DatasourceGetPagesResponse: - del datasource_parameters - return DatasourceGetPagesResponse(result=[]) - - def _get_content( - self, page: GetOnlineDocumentPageContentRequest - ) -> Generator[DatasourceMessage, None, None]: - yield DatasourceMessage( - type=DatasourceMessage.MessageType.TEXT, - message=DatasourceMessage.TextMessage( - text=( - "Online document page content for " - f"{page.workspace_id} - {page.page_id}" - ) - ), - ) - - session = Session( - session_id="test", - executor=ThreadPoolExecutor(max_workers=1), - reader=StdioRequestReader(), - writer=StdioResponseWriter(), - ) - datasource = OnlineDocument( - runtime=DatasourceRuntime(credentials={}, user_id="test", session_id="test"), - session=session, - ) - assert datasource is not None diff --git a/tests/interfaces/datasource/test_datasource_workflows.py b/tests/interfaces/datasource/test_datasource_workflows.py new file mode 100644 index 00000000..23b594a2 --- /dev/null +++ b/tests/interfaces/datasource/test_datasource_workflows.py @@ -0,0 +1,185 @@ +import pytest + +from dify_plugin.core.runtime import Session +from dify_plugin.entities.datasource import ( + DatasourceMessage, + DatasourceRuntime, + GetOnlineDocumentPageContentRequest, + OnlineDriveBrowseFilesRequest, + OnlineDriveDownloadFileRequest, + WebsiteCrawlMessage, +) +from dify_plugin.plugin import Plugin +from tests.core.plugin_fixtures import contract_plugin, contract_source + +__all__ = ["contract_plugin", "contract_source"] + + +@pytest.fixture +def session(contract_plugin: Plugin) -> Session: + return Session( + session_id="datasource-session", + executor=contract_plugin.executer, + reader=contract_plugin.request_reader, + writer=contract_plugin.default_response_writer, + ) + + +@pytest.fixture +def runtime() -> DatasourceRuntime: + return DatasourceRuntime( + credentials={"api_key": "valid"}, + user_id="user", + session_id="datasource-session", + ) + + +@pytest.mark.parametrize("outcome", ["complete", "close", "error"]) +def test_website_crawl_stream_status_errors_and_cleanup( + contract_plugin: Plugin, session: Session, runtime: DatasourceRuntime, outcome: str +) -> None: + datasource = contract_plugin.registration.get_website_crawl_datasource_cls( + "website", "echo" + )(runtime=runtime, session=session) + stream = datasource.website_crawl({ + "url": "https://example.invalid/", + "fail": outcome == "error", + }) + assert not hasattr(datasource, "finished") + first = next(stream) + assert isinstance(first, WebsiteCrawlMessage) + assert first.result.status == "processing" + assert first.result.completed == 0 + assert not datasource.finished + if outcome == "close": + stream.close() + elif outcome == "error": + with pytest.raises(ValueError, match="crawl failed"): + next(stream) + else: + last = next(stream) + assert last.result.status == "completed" + assert last.result.completed == last.result.total == 1 + assert last.result.web_info_list[0].source_url == "https://example.invalid/" + assert last.result.web_info_list[0].content == "page content" + with pytest.raises(StopIteration): + next(stream) + assert datasource.finished + assert datasource.runtime.credentials == {"api_key": "valid"} + assert datasource.session is session + + +def test_document_listing_drives_content_request_and_empty_result( + contract_plugin: Plugin, session: Session, runtime: DatasourceRuntime +) -> None: + datasource = contract_plugin.registration.get_online_document_datasource_cls( + "documents", "echo" + )(runtime=runtime, session=session) + listing = datasource.get_pages({}) + assert len(listing.result) == 1 + workspace = listing.result[0] + assert workspace.total == len(workspace.pages) == 1 + page = workspace.pages[0] + stream = datasource.get_content( + GetOnlineDocumentPageContentRequest( + workspace_id=workspace.workspace_id, page_id=page.page_id, type=page.type + ) + ) + assert not hasattr(datasource, "finished") + messages = list(stream) + assert all(isinstance(message, DatasourceMessage) for message in messages) + assert all( + message.type == DatasourceMessage.MessageType.VARIABLE for message in messages + ) + assert { + message.message.variable_name: message.message.variable_value + for message in messages + } == {"content": "page content", "page_id": "page", "workspace_id": "workspace"} + assert datasource.finished + assert datasource.get_pages({"empty": True}).result == [] + + +@pytest.mark.parametrize("outcome", ["close", "error"]) +def test_document_content_early_close_and_missing_page( + contract_plugin: Plugin, session: Session, runtime: DatasourceRuntime, outcome: str +) -> None: + datasource = contract_plugin.registration.get_online_document_datasource_cls( + "documents", "echo" + )(runtime=runtime, session=session) + stream = datasource.get_content( + GetOnlineDocumentPageContentRequest( + workspace_id="workspace", + page_id="missing" if outcome == "error" else "page", + type="page", + ) + ) + if outcome == "error": + with pytest.raises(ValueError, match="page not found"): + next(stream) + else: + assert next(stream).message.variable_name == "content" + assert not datasource.finished + stream.close() + assert datasource.finished + + +def test_drive_pagination_keeps_cursor_bucket_and_remote_metadata( + contract_plugin: Plugin, session: Session, runtime: DatasourceRuntime +) -> None: + datasource = contract_plugin.registration.get_online_drive_datasource_cls( + "drive", "echo" + )(runtime=runtime, session=session) + first = datasource.browse_files( + OnlineDriveBrowseFilesRequest(bucket="bucket", prefix="folder", max_keys=1) + ).result[0] + assert first.bucket == "bucket" + assert first.is_truncated + assert first.next_page_parameters == {"cursor": "next"} + assert first.files[0].id == "file" + assert first.files[0].remote_metadata == {"etag": '"v1"'} + last = datasource.browse_files( + OnlineDriveBrowseFilesRequest( + bucket=first.bucket, + prefix="folder", + max_keys=1, + next_page_parameters=first.next_page_parameters, + ) + ).result[0] + assert last.bucket == first.bucket + assert last.files[0].id == "last-file" + assert not last.is_truncated + assert last.next_page_parameters is None + assert ( + datasource.browse_files(OnlineDriveBrowseFilesRequest(prefix="empty")).result + == [] + ) + + +@pytest.mark.parametrize("outcome", ["complete", "close", "error"]) +def test_drive_download_preserves_binary_metadata_errors_and_cleanup( + contract_plugin: Plugin, session: Session, runtime: DatasourceRuntime, outcome: str +) -> None: + datasource = contract_plugin.registration.get_online_drive_datasource_cls( + "drive", "echo" + )(runtime=runtime, session=session) + stream = datasource.download_file( + OnlineDriveDownloadFileRequest( + bucket="bucket", id="broken" if outcome == "error" else "file" + ) + ) + assert not hasattr(datasource, "finished") + message = next(stream) + assert isinstance(message, DatasourceMessage) + assert message.type == DatasourceMessage.MessageType.BLOB + assert message.message.blob == b"data" + assert message.meta == {"mime_type": "text/plain", "file_name": "example.txt"} + assert not datasource.finished + if outcome == "close": + stream.close() + elif outcome == "error": + with pytest.raises(ValueError, match="download interrupted"): + next(stream) + else: + with pytest.raises(StopIteration): + next(stream) + assert datasource.finished diff --git a/tests/interfaces/model/openai_compatible/conftest.py b/tests/interfaces/model/openai_compatible/conftest.py new file mode 100644 index 00000000..675f4fb5 --- /dev/null +++ b/tests/interfaces/model/openai_compatible/conftest.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from gevent.event import Event + +from dify_plugin.interfaces.model.openai_compatible import llm + +if TYPE_CHECKING: + from requests import Response + + +@pytest.fixture +def http_responses(monkeypatch: pytest.MonkeyPatch) -> tuple[list[Response], Event]: + responses = [] + ready = Event() + post = llm.requests.post + + def observe_response(url: str, **kwargs: object) -> Response: + response = post(url, **kwargs) + responses.append(response) + ready.set() + return response + + monkeypatch.setattr(llm.requests, "post", observe_response) + return responses, ready diff --git a/tests/interfaces/model/openai_compatible/test_http_audio.py b/tests/interfaces/model/openai_compatible/test_http_audio.py new file mode 100644 index 00000000..204ae0a0 --- /dev/null +++ b/tests/interfaces/model/openai_compatible/test_http_audio.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +import json +import wave +from email import policy +from email.parser import BytesParser +from io import BytesIO +from pathlib import Path +from typing import TYPE_CHECKING + +import gevent +import pytest +from gevent.event import Event + +from dify_plugin.errors.model import CredentialsValidateFailedError, InvokeError +from dify_plugin.interfaces.model.openai_compatible.speech2text import ( + OAICompatSpeech2TextModel, +) +from dify_plugin.interfaces.model.openai_compatible.tts import OAICompatText2SpeechModel +from tests.http_server import HTTPServerFactory, Response + +if TYPE_CHECKING: + from requests import Response as HTTPResponse + +pytestmark = pytest.mark.integration + + +def wav_audio(frames: bytes, rate: int = 24_000) -> bytes: + output = BytesIO() + with wave.open(output, "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(rate) + audio.writeframes(frames) + return output.getvalue() + + +@pytest.mark.parametrize("count", [1, 2]) +def test_tts_preserves_provider_wav_bytes_for_each_sentence( + http_server: HTTPServerFactory, + count: int, +) -> None: + frames = [b"\x01\x00" * 8, b"\x02\x00" * 8][:count] + audio = [wav_audio(frame) for frame in frames] + sentences = ["One.", "Two."][:count] + server = http_server([Response(chunks=(body[:13], body[13:])) for body in audio]) + model = OAICompatText2SpeechModel([]) + + result = b"".join( + model.invoke( + "model", + "tenant", + { + "endpoint_url": f"{server.url}/v1/", + "api_key": "key", + "endpoint_model_name": "upstream-tts", + "word_limit": 4, + "audio_type": "wav", + }, + "".join(sentences), + "voice-1", + ) + ) + + assert result == b"".join(audio) + assert [request.path for request in server.requests] == ["/v1/audio/speech"] * count + assert all( + request.headers["Authorization"] == "Bearer key" for request in server.requests + ) + assert [json.loads(request.body) for request in server.requests] == [ + { + "model": "upstream-tts", + "input": sentence, + "voice": "voice-1", + "response_format": "wav", + } + for sentence in sentences + ] + assert model.started_at == 0 + + +@pytest.mark.parametrize("failure", ["format", "truncated", "invalid"]) +def test_tts_streams_each_provider_segment_without_reencoding( + http_server: HTTPServerFactory, + failure: str, +) -> None: + valid = wav_audio(b"\x01\x00" * 8) + second = { + "format": wav_audio(b"\x02\x00" * 8, rate=16_000), + "truncated": valid[:-2], + "invalid": b"not WAV", + }[failure] + server = http_server([Response(body=valid), Response(body=second)]) + model = OAICompatText2SpeechModel([]) + stream = model.invoke( + "model", + "tenant", + {"endpoint_url": server.url, "word_limit": 4, "audio_type": "wav"}, + "One.Two.", + "alloy", + ) + try: + assert next(stream) == valid + assert len(server.requests) == 1 + assert b"".join(stream) == second + finally: + stream.close() + assert model.started_at == 0 + + +def test_tts_empty_text_does_not_make_a_request(http_server: HTTPServerFactory) -> None: + server = http_server([]) + chunks = list( + OAICompatText2SpeechModel([]).invoke( + "model", "tenant", {"endpoint_url": server.url}, "", "alloy" + ) + ) + assert chunks == [] + assert server.requests == [] + + +def test_tts_early_close_releases_the_http_response( + http_server: HTTPServerFactory, http_responses: tuple[list[HTTPResponse], Event] +) -> None: + responses, _ready = http_responses + server = http_server([Response(chunks=(b"first", b"second"))]) + model = OAICompatText2SpeechModel([]) + stream = model.invoke( + "model", "tenant", {"endpoint_url": server.url}, "Hello", "alloy" + ) + + assert next(stream) == b"first" + stream.close() + + assert responses[0].raw.closed + assert model.started_at == 0 + + +def test_tts_greenlet_cancellation_releases_the_http_response( + http_server: HTTPServerFactory, http_responses: tuple[list[HTTPResponse], Event] +) -> None: + responses, ready = http_responses + server = http_server([Response(chunks=(b"delayed",), delay=10)]) + model = OAICompatText2SpeechModel([]) + stream = model.invoke( + "model", "tenant", {"endpoint_url": server.url}, "Hello", "alloy" + ) + worker = gevent.spawn(list, stream) + try: + assert ready.wait(timeout=2), "The HTTP response was not received" + finally: + worker.kill(timeout=2) + + assert worker.ready() + assert responses[0].raw.closed + assert model.started_at == 0 + + +def test_tts_truncated_response_reports_error_and_closes_response( + http_server: HTTPServerFactory, http_responses: tuple[list[HTTPResponse], Event] +) -> None: + responses, _ready = http_responses + server = http_server([ + Response(headers={"Content-Length": "5000"}, body=b"a" * 4096) + ]) + model = OAICompatText2SpeechModel([]) + stream = model.invoke( + "model", "tenant", {"endpoint_url": server.url}, "Hello", "alloy" + ) + + chunks = [] + with pytest.raises(InvokeError): + chunks.extend(stream) + + assert (b"a" * 4096).startswith(b"".join(chunks)) + assert responses[0].raw.closed + assert model.started_at == 0 + + +def test_tts_credential_validation_closes_the_sample_stream( + http_server: HTTPServerFactory, + http_responses: tuple[list[HTTPResponse], Event], + monkeypatch: pytest.MonkeyPatch, +) -> None: + responses, _ready = http_responses + server = http_server([Response(chunks=(b"sample", b"unused audio"))]) + model = OAICompatText2SpeechModel([]) + credentials = {"endpoint_url": server.url, "voices": "voice-1,voice-2"} + stream = model._invoke("model", "validate", credentials, "Test.", "voice-1") + monkeypatch.setattr(model, "_invoke", lambda **_kwargs: stream) + + model.validate_credentials("model", credentials) + + assert responses[0].raw.closed + assert list(stream) == [] + assert json.loads(server.requests[0].body)["voice"] == "voice-1" + + +@pytest.mark.parametrize("validate", [False, True]) +def test_tts_http_failure_closes_the_response( + http_server: HTTPServerFactory, + http_responses: tuple[list[HTTPResponse], Event], + validate: bool, +) -> None: + responses, _ready = http_responses + server = http_server([Response(status=401, body=b"invalid API key")]) + model = OAICompatText2SpeechModel([]) + credentials = {"endpoint_url": server.url} + if validate: + with pytest.raises(CredentialsValidateFailedError, match="invalid API key"): + model.validate_credentials("model", credentials) + else: + with pytest.raises(InvokeError, match="invalid API key"): + list(model.invoke("model", "tenant", credentials, "Hello", "alloy")) + + assert responses[0].raw.closed + assert model.started_at == 0 + + +@pytest.mark.parametrize( + ("voices", "expected"), + [ + (" alloy, , echo ", ["alloy", "echo"]), + (" , ", ["alloy"]), + ], +) +def test_tts_voice_choices_match_credential_configuration( + voices: str, expected: list[str] +) -> None: + model = OAICompatText2SpeechModel([]) + choices = model.get_tts_model_voices("model", {"voices": voices}) + assert [choice["value"] for choice in choices] == expected + + +def test_speech_to_text_sends_multipart_without_closing_the_callers_file( + http_server: HTTPServerFactory, tmp_path: Path +) -> None: + audio = b"RIFF\x00\xffWAVEaudio" + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(audio) + server = http_server([Response(body='{"text":"你好"}'.encode())]) + model = OAICompatSpeech2TextModel([]) + + with audio_path.open("rb") as audio_file: + result = model.invoke( + "model", + { + "endpoint_url": f"{server.url}/v1/", + "api_key": "key", + "endpoint_model_name": "upstream-stt", + }, + audio_file, + ) + assert not audio_file.closed + + assert result == "你好" + assert model.started_at == 0 + request = server.requests[0] + assert request.path == "/v1/audio/transcriptions" + assert request.headers["Authorization"] == "Bearer key" + multipart = BytesParser(policy=policy.default).parsebytes( + b"Content-Type: " + + request.headers["Content-Type"].encode() + + b"\r\n\r\n" + + request.body + ) + fields = { + part.get_param("name", header="content-disposition"): part + for part in multipart.iter_parts() + } + assert fields["model"].get_payload(decode=True) == b"upstream-stt" + assert fields["file"].get_filename() == "audio.wav" + assert fields["file"].get_payload(decode=True) == audio + + +@pytest.mark.parametrize( + "response", + [ + Response(status=503, body=b"upstream unavailable"), + Response(body=b"not json"), + Response(body=b"{}"), + ], +) +def test_speech_to_text_errors_preserve_file_ownership_and_timing( + http_server: HTTPServerFactory, tmp_path: Path, response: Response +) -> None: + audio_path = tmp_path / "audio.wav" + audio_path.write_bytes(b"audio") + server = http_server([response]) + model = OAICompatSpeech2TextModel([]) + with audio_path.open("rb") as audio_file: + with pytest.raises(InvokeError): + model.invoke("model", {"endpoint_url": server.url}, audio_file) + assert not audio_file.closed + assert model.started_at == 0 + + +@pytest.mark.parametrize("status", [200, 401]) +def test_speech_to_text_validates_credentials_with_sample_audio( + http_server: HTTPServerFactory, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + status: int, +) -> None: + audio_path = tmp_path / "sample.wav" + audio_path.write_bytes(b"sample audio") + server = http_server([Response(status=status, body=b'{"text":"sample"}')]) + model = OAICompatSpeech2TextModel([]) + monkeypatch.setattr(model, "_get_demo_file_path", lambda: str(audio_path)) + if status == 200: + model.validate_credentials("model", {"endpoint_url": server.url}) + else: + with pytest.raises(CredentialsValidateFailedError): + model.validate_credentials("model", {"endpoint_url": server.url}) + assert b"sample audio" in server.requests[0].body diff --git a/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py b/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py new file mode 100644 index 00000000..b6768c94 --- /dev/null +++ b/tests/interfaces/model/openai_compatible/test_http_embedding_rerank.py @@ -0,0 +1,465 @@ +import json +import math +from decimal import Decimal + +import pytest + +from dify_plugin.entities import I18nObject +from dify_plugin.entities.model import ( + AIModelEntity, + ModelPropertyKey, + ModelType, + PriceConfig, +) +from dify_plugin.errors.model import CredentialsValidateFailedError, InvokeError +from dify_plugin.interfaces.model.openai_compatible.rerank import OAICompatRerankModel +from dify_plugin.interfaces.model.openai_compatible.text_embedding import ( + OAICompatEmbeddingModel, +) +from tests.http_server import HTTPServerFactory, Response + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def embedding(monkeypatch: pytest.MonkeyPatch) -> OAICompatEmbeddingModel: + model = OAICompatEmbeddingModel([ + AIModelEntity( + model="model", + label=I18nObject(en_us="model"), + model_type=ModelType.TEXT_EMBEDDING, + model_properties={ + ModelPropertyKey.CONTEXT_SIZE: 5, + ModelPropertyKey.MAX_CHUNKS: 2, + }, + parameter_rules=[], + pricing=PriceConfig(input=Decimal(2), unit=Decimal("0.01"), currency="USD"), + ) + ]) + monkeypatch.setattr(model, "_get_num_tokens_by_gpt2", len) + return model + + +def test_embedding_batches_truncates_and_accumulates_usage( + http_server: HTTPServerFactory, embedding: OAICompatEmbeddingModel +) -> None: + server = http_server([ + Response( + body=b'{"data":[{"embedding":[1,2]},{"embedding":[3,4]}],"usage":{"total_tokens":5}}' + ), + Response(body=b'{"data":[{"embedding":[5,6]}],"usage":{"total_tokens":2}}'), + ]) + result = embedding.invoke( + "model", + { + "endpoint_url": f"{server.url}/v1/", + "endpoint_model_name": "upstream-model", + "api_key": "key", + }, + ["abcdef", "你好", ""], + user="user-1", + ) + + assert result.embeddings == [[1, 2], [3, 4], [5, 6]] + assert result.model == "model" + assert result.usage.total_tokens == 7 + assert result.usage.total_price == Decimal("0.14") + assert result.usage.latency >= 0 + assert embedding.started_at == 0 + requests = server.requests + assert [request.path for request in requests] == [ + "/v1/embeddings", + "/v1/embeddings", + ] + assert all(request.headers["Authorization"] == "Bearer key" for request in requests) + payloads = [json.loads(request.body) for request in requests] + assert [payload["input"] for payload in payloads] == [["abcde", "你好"], [""]] + assert all(payload["model"] == "upstream-model" for payload in payloads) + assert all(payload["encoding_format"] == "float" for payload in payloads) + assert all(payload["user"] == "user-1" for payload in payloads) + + +def test_embedding_empty_input_does_not_request_the_server( + http_server: HTTPServerFactory, embedding: OAICompatEmbeddingModel +) -> None: + server = http_server([]) + result = embedding.invoke("model", {"endpoint_url": server.url}, []) + assert result.embeddings == [] + assert result.usage.total_tokens == 0 + assert server.requests == [] + + +@pytest.mark.parametrize( + "response", + [ + Response(status=429, body=b"rate limited"), + Response(body=b"not json"), + ], +) +def test_embedding_failure_stops_later_batches( + http_server: HTTPServerFactory, + embedding: OAICompatEmbeddingModel, + response: Response, +) -> None: + server = http_server([ + response, + Response(body=b'{"data":[{"embedding":[5,6]}],"usage":{"total_tokens":1}}'), + ]) + try: + with pytest.raises(InvokeError): + embedding.invoke( + "model", {"endpoint_url": server.url}, ["one", "two", "three"] + ) + finally: + assert len(server.requests) == 1 + assert embedding.started_at == 0 + + +@pytest.mark.parametrize("reordered", [False, True]) +def test_embedding_preserves_provider_vector_order( + http_server: HTTPServerFactory, + embedding: OAICompatEmbeddingModel, + reordered: bool, +) -> None: + batch = [{"index": 0, "embedding": [1, 2]}, {"index": 1, "embedding": [3, 4]}] + if reordered: + batch.reverse() + server = http_server([ + Response( + body=json.dumps({"data": batch, "usage": {"total_tokens": 2}}).encode() + ), + Response( + body=b'{"data":[{"index":0,"embedding":[5,6]}],"usage":{"total_tokens":1}}' + ), + ]) + + result = embedding.invoke( + "model", {"endpoint_url": server.url}, ["one", "two", "three"] + ) + + expected = [[3, 4], [1, 2]] if reordered else [[1, 2], [3, 4]] + assert result.embeddings == [*expected, [5, 6]] + assert result.usage.total_tokens == 3 + assert [json.loads(request.body)["input"] for request in server.requests] == [ + ["one", "two"], + ["three"], + ] + assert embedding.started_at == 0 + + +@pytest.mark.parametrize( + ("data_json", "expected"), + [ + ("[]", []), + ('[{"index":0,"embedding":[1,2]}]', [[1, 2]]), + ( + '[{"index":0,"embedding":[1,2]},{"index":0,"embedding":[3,4]}]', + [[1, 2], [3, 4]], + ), + ( + '[{"index":0,"embedding":[1,2]},{"index":2,"embedding":[3,4]}]', + [[1, 2], [3, 4]], + ), + ('[{"index":0,"embedding":[]},{"index":1,"embedding":[3,4]}]', [[], [3, 4]]), + ( + '[{"index":0,"embedding":[1e400,1]},{"index":1,"embedding":[3,4]}]', + [[float("inf"), 1], [3, 4]], + ), + ('[{"index":true,"embedding":[true,"2.5"]}]', [[1, 2.5]]), + ], + ids=[ + "empty-batch", + "missing-vector", + "duplicate-index", + "out-of-range-index", + "empty-vector", + "non-finite-vector", + "coercible-vector", + ], +) +def test_embedding_keeps_existing_response_acceptance( + http_server: HTTPServerFactory, + embedding: OAICompatEmbeddingModel, + data_json: str, + expected: list[list[float]], +) -> None: + server = http_server([ + Response(body=f'{{"data":{data_json},"usage":{{"total_tokens":2}}}}'.encode()), + Response(body=b'{"data":[{"embedding":[5,6]}],"usage":{"total_tokens":1}}'), + ]) + + result = embedding.invoke( + "model", {"endpoint_url": server.url}, ["one", "two", "three"] + ) + assert result.embeddings == [*expected, [5, 6]] + assert result.usage.total_tokens == 3 + assert len(server.requests) == 2 + assert embedding.started_at == 0 + + +def test_embedding_uses_safe_batch_defaults_when_schema_omits_limits( + http_server: HTTPServerFactory, +) -> None: + server = http_server([ + Response(body=b'{"data":[{"embedding":[1,2]}],"usage":{"total_tokens":1}}'), + Response(body=b'{"data":[{"embedding":[3,4]}],"usage":{"total_tokens":1}}'), + ]) + model = OAICompatEmbeddingModel([ + AIModelEntity( + model="model", model_type=ModelType.TEXT_EMBEDDING, model_properties={} + ) + ]) + + result = model.invoke("model", {"endpoint_url": server.url}, ["Hello", "World"]) + + assert result.embeddings == [[1, 2], [3, 4]] + assert result.usage.total_tokens == 2 + assert [json.loads(request.body)["input"] for request in server.requests] == [ + ["Hello"], + ["World"], + ] + + +@pytest.mark.parametrize( + ("response", "valid"), + [ + (Response(body=b'{"model":"upstream"}'), True), + (Response(status=401, body=b"unauthorized"), False), + (Response(body=b"not json"), False), + (Response(body=b"{}"), False), + (Response(body=b"1"), False), + (Response(body=b'"model"'), True), + (Response(body=b'["model"]'), True), + ], +) +def test_embedding_validates_credentials_over_http( + http_server: HTTPServerFactory, response: Response, valid: bool +) -> None: + server = http_server([response]) + model = OAICompatEmbeddingModel([]) + credentials = {"endpoint_url": server.url, "api_key": "key"} + if valid: + model.validate_credentials("model", credentials) + else: + with pytest.raises(CredentialsValidateFailedError): + model.validate_credentials("model", credentials) + request = server.requests[0] + assert request.path == "/embeddings" + assert request.headers["Authorization"] == "Bearer key" + assert json.loads(request.body) == {"input": ["ping"], "model": "model"} + + +@pytest.mark.parametrize("threshold", [None, 0.5]) +def test_rerank_uses_response_indices_documents_and_normalized_scores( + http_server: HTTPServerFactory, threshold: float | None +) -> None: + server = http_server([ + Response( + body=json.dumps({ + "results": [ + { + "index": 2, + "relevance_score": -2, + "document": {"text": "replacement"}, + }, + {"index": 0, "relevance_score": -4, "document": "original"}, + {"index": 1, "relevance_score": 0}, + ] + }).encode() + ) + ]) + model = OAICompatRerankModel([]) + result = model.invoke( + "model", + {"endpoint_url": server.url, "endpoint_model_name": "reranker"}, + "query", + ["zero", "one", "two"], + score_threshold=threshold, + top_n=3, + ) + + expected = [(1, "one", 1), (2, "replacement", 0.5)] + if threshold is None: + expected.append((0, "original", 0)) + assert [(doc.index, doc.text, doc.score) for doc in result.docs] == expected + request = server.requests[0] + assert request.path == "/rerank" + assert "Authorization" not in request.headers + assert json.loads(request.body) == { + "model": "reranker", + "query": "query", + "documents": ["zero", "one", "two"], + "top_n": 3, + "return_documents": True, + } + assert model.started_at == 0 + + +def test_rerank_keeps_empty_response_error(http_server: HTTPServerFactory) -> None: + server = http_server([Response(body=b'{"results":[]}')]) + with pytest.raises(InvokeError, match="min\\(\\)"): + OAICompatRerankModel([]).invoke( + "model", {"endpoint_url": server.url}, "query", ["document"] + ) + + +def test_rerank_empty_documents_does_not_request_the_server( + http_server: HTTPServerFactory, +) -> None: + server = http_server([]) + result = OAICompatRerankModel([]).invoke( + "model", {"endpoint_url": server.url}, "query", [] + ) + assert result.docs == [] + assert server.requests == [] + + +def test_rerank_equal_scores_remain_finite(http_server: HTTPServerFactory) -> None: + server = http_server([ + Response( + body=b'{"results":[{"index":0,"relevance_score":2},{"index":1,"relevance_score":2}]}' + ) + ]) + result = OAICompatRerankModel([]).invoke( + "model", {"endpoint_url": server.url}, "query", ["first", "second"] + ) + assert [(doc.text, doc.score) for doc in result.docs] == [ + ("first", 0), + ("second", 0), + ] + + +@pytest.mark.parametrize( + ("scores", "threshold", "expected"), + [ + ([0.9], 0.5, []), + ([0.8, 0.9], 0.5, [(1, 1.0)]), + ([0.8, 0.9], 0.9, [(1, 1.0)]), + ([0.9, 0.9], 0.5, []), + ([0.0, 1.0], None, [(1, 1.0), (0, 0.0)]), + ], +) +def test_rerank_preserves_batch_normalization_and_thresholds( + http_server: HTTPServerFactory, + scores: list[float], + threshold: float | None, + expected: list[tuple[int, float]], +) -> None: + server = http_server([ + Response( + body=json.dumps({ + "results": [ + {"index": index, "relevance_score": score} + for index, score in enumerate(scores) + ] + }).encode() + ) + ]) + + result = OAICompatRerankModel([]).invoke( + "model", + {"endpoint_url": server.url}, + "query", + ["first", "second"][: len(scores)], + score_threshold=threshold, + ) + + assert [(doc.index, doc.score) for doc in result.docs] == expected + + +@pytest.mark.parametrize( + "response", + [ + Response(status=503, body=b"unavailable"), + Response(body=b"invalid json"), + Response(body=b'{"results":[{"index":2,"relevance_score":1}]}'), + Response(body=b'{"results":[{"index":0,"relevance_score":"0.9"}]}'), + ], +) +def test_rerank_errors_reach_the_public_interface( + http_server: HTTPServerFactory, response: Response +) -> None: + server = http_server([response]) + model = OAICompatRerankModel([]) + with pytest.raises(InvokeError): + model.invoke("model", {"endpoint_url": server.url}, "query", ["document"]) + assert model.started_at == 0 + + +@pytest.mark.parametrize( + ("results_json", "expected_indices", "non_finite"), + [ + ('[{"index":-1,"relevance_score":0.9}]', [-1], [False]), + ('[{"index":true,"relevance_score":0.9}]', [1], [False]), + ('[{"index":0,"relevance_score":true}]', [0], [False]), + ('[{"index":0,"relevance_score":1e400}]', [0], [True]), + ('[{"index":0,"relevance_score":-1e400}]', [0], [True]), + ( + '[{"index":0,"relevance_score":1e308},{"index":1,"relevance_score":-1e308}]', + [0, 1], + [True, False], + ), + ], + ids=[ + "negative-index", + "boolean-index", + "boolean-score", + "positive-overflow", + "negative-overflow", + "normalization-overflow", + ], +) +def test_rerank_keeps_existing_response_acceptance( + http_server: HTTPServerFactory, + results_json: str, + expected_indices: list[int], + non_finite: list[bool], +) -> None: + body = f'{{"results":{results_json}}}'.encode() + server = http_server([Response(body=body)]) + model = OAICompatRerankModel([]) + + result = model.invoke( + "model", {"endpoint_url": server.url}, "query", ["first", "second"] + ) + assert [doc.index for doc in result.docs] == expected_indices + assert [doc.text for doc in result.docs] == [ + ["first", "second"][index] for index in expected_indices + ] + assert [math.isnan(doc.score) for doc in result.docs] == non_finite + assert all(doc.score == 0 for doc in result.docs if not math.isnan(doc.score)) + assert len(server.requests) == 1 + assert model.started_at == 0 + + +@pytest.mark.parametrize("status", [200, 401]) +def test_rerank_validates_credentials_over_http( + http_server: HTTPServerFactory, status: int +) -> None: + server = http_server([ + Response(status=status, body=b'{"results":[{"index":0,"relevance_score":1}]}') + ]) + model = OAICompatRerankModel([]) + if status == 200: + model.validate_credentials("model", {"endpoint_url": server.url}) + else: + with pytest.raises(CredentialsValidateFailedError): + model.validate_credentials("model", {"endpoint_url": server.url}) + assert json.loads(server.requests[0].body)["query"] + + +@pytest.mark.parametrize( + ("model_name", "endpoint_url", "error"), + [("model", "", "server_url is required"), ("", None, "model_name is required")], +) +def test_rerank_rejects_missing_configuration_before_requesting( + http_server: HTTPServerFactory, + model_name: str, + endpoint_url: str | None, + error: str, +) -> None: + server = http_server([]) + credentials = {"endpoint_url": server.url if endpoint_url is None else endpoint_url} + with pytest.raises(CredentialsValidateFailedError, match=error): + OAICompatRerankModel([]).validate_credentials(model_name, credentials) + assert server.requests == [] diff --git a/tests/interfaces/model/openai_compatible/test_http_llm.py b/tests/interfaces/model/openai_compatible/test_http_llm.py new file mode 100644 index 00000000..ab0038b8 --- /dev/null +++ b/tests/interfaces/model/openai_compatible/test_http_llm.py @@ -0,0 +1,940 @@ +from __future__ import annotations + +import json +from collections.abc import Generator +from decimal import Decimal +from typing import TYPE_CHECKING + +import gevent +import pytest +from gevent.event import Event + +from dify_plugin.entities import I18nObject +from dify_plugin.entities.model import ParameterRule, ParameterType +from dify_plugin.entities.model.llm import LLMResult, LLMResultChunk +from dify_plugin.entities.model.message import ( + AssistantPromptMessage, + PromptMessage, + PromptMessageTool, + UserPromptMessage, +) +from dify_plugin.errors.model import CredentialsValidateFailedError, InvokeError +from dify_plugin.interfaces.model.openai_compatible import llm +from dify_plugin.interfaces.model.openai_compatible.llm import ( + OAICompatLargeLanguageModel, +) +from tests.http_server import HTTPServerFactory, Response + +if TYPE_CHECKING: + from requests import Response as HTTPResponse + +pytestmark = pytest.mark.integration + + +class LegacyResponseLLM(OAICompatLargeLanguageModel): + """Use the response hooks implemented by existing official plugins.""" + + def _handle_generate_response( + self, + model: str, + credentials: dict, + response: HTTPResponse, + prompt_messages: list[PromptMessage], + ) -> LLMResult: + return super()._handle_generate_response( + model, credentials, response, prompt_messages + ) + + def _handle_generate_stream_response( + self, + model: str, + credentials: dict, + response: HTTPResponse, + prompt_messages: list[PromptMessage], + ) -> Generator[LLMResultChunk]: + yield from super()._handle_generate_stream_response( + model, credentials, response, prompt_messages + ) + + def _create_final_llm_result_chunk( + self, + index: int, + message: AssistantPromptMessage, + finish_reason: str | None, + usage: dict | None, + model: str, + prompt_messages: list[PromptMessage], + credentials: dict, + full_content: str, + ) -> LLMResultChunk: + return super()._create_final_llm_result_chunk( + index, + message, + finish_reason, + usage, + model, + prompt_messages, + credentials, + full_content, + ) + + +@pytest.fixture +def structured_model() -> OAICompatLargeLanguageModel: + schema = OAICompatLargeLanguageModel([]).get_customizable_model_schema( + "model", {"mode": "chat"} + ) + # Official compatible providers expose these options through their schema. + schema.parameter_rules.extend([ + ParameterRule( + name="response_format", + label=I18nObject(en_us="Response format"), + type=ParameterType.STRING, + options=["text", "json_object", "json_schema"], + ), + ParameterRule( + name="json_schema", + label=I18nObject(en_us="JSON schema"), + type=ParameterType.STRING, + ), + ]) + return OAICompatLargeLanguageModel([schema]) + + +@pytest.mark.parametrize( + ("parameters", "expected"), + [ + ( + { + "response_format": "json_schema", + "json_schema": '{"name":"answer","schema":{"type":"object"}}', + }, + { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": {"type": "object"}}, + }, + ), + ({"response_format": "json_object"}, {"type": "json_object"}), + ({"response_format": "text"}, {"type": "text"}), + ({"json_schema": "unused"}, None), + ], +) +def test_structured_output_parameters_reach_the_provider( + http_server: HTTPServerFactory, + structured_model: OAICompatLargeLanguageModel, + parameters: dict, + expected: dict | None, +) -> None: + server = http_server([ + Response( + body=b'{"choices":[{"message":{"content":"{}"}}],' + b'"usage":{"prompt_tokens":1,"completion_tokens":1}}' + ) + ]) + + chunks = list( + structured_model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Answer as JSON")], + model_parameters=parameters, + stream=False, + ) + ) + + assert chunks[-1].delta.message.content == "{}" + payload = json.loads(server.requests[0].body) + assert payload.get("response_format") == expected + assert "json_schema" not in payload + + +@pytest.mark.parametrize("schema", [None, "", "{broken", "[]"]) +def test_invalid_structured_schema_fails_before_sending_a_request( + http_server: HTTPServerFactory, + structured_model: OAICompatLargeLanguageModel, + schema: str | None, +) -> None: + server = http_server([]) + with pytest.raises(InvokeError, match=r"[Ss]chema"): + list( + structured_model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Answer as JSON")], + model_parameters={ + "response_format": "json_schema", + "json_schema": schema, + }, + stream=False, + ) + ) + assert server.requests == [] + assert structured_model.started_at == 0 + + +@pytest.mark.parametrize("mode", ["chat", "completion"]) +def test_invoke_sends_json_and_reads_non_stream_response( + http_server: HTTPServerFactory, mode: str +) -> None: + choice = {"message": {"content": "回答"}} if mode == "chat" else {"text": "回答"} + server = http_server([ + Response( + body=json.dumps({ + "id": "response-1", + "model": "upstream-model", + "choices": [choice], + "usage": {"prompt_tokens": 4, "completion_tokens": 2}, + }).encode() + ) + ]) + content = "你好😀\ud800" + chunks = list( + OAICompatLargeLanguageModel([]).invoke( + model="display-model", + credentials={ + "mode": mode, + "endpoint_url": f"{server.url}/v1/", + "endpoint_model_name": "request-model", + "api_key": "test-key", + "extra_headers": { + "Content-Type": "application/custom+json", + "X-Dify-App-Id": "app-1", + }, + }, + prompt_messages=[UserPromptMessage(content=content)], + model_parameters={"temperature": 0.25}, + stop=["stop"], + user="user-1", + stream=False, + ) + ) + + assert len(chunks) == 1 + assert chunks[0].model == "upstream-model" + assert chunks[0].delta.message.content == "回答" + assert chunks[0].delta.usage is not None + assert chunks[0].delta.usage.total_tokens == 6 + request = server.requests[0] + assert request.method == "POST" + assert request.path == ( + "/v1/chat/completions" if mode == "chat" else "/v1/completions" + ) + assert request.headers["Authorization"] == "Bearer test-key" + assert request.headers["X-Dify-App-Id"] == "app-1" + assert request.headers["Content-Type"] == "application/custom+json" + assert "你好😀".encode() in request.body + assert b"\\ud800" in request.body + payload = json.loads(request.body) + assert payload["model"] == "request-model" + assert payload["stream"] is False + assert payload["temperature"] == pytest.approx(0.25) + assert payload["stop"] == ["stop"] + assert payload["user"] == "user-1" + if mode == "chat": + assert payload["messages"] == [{"role": "user", "content": content}] + else: + assert payload["prompt"] == content + + +@pytest.mark.parametrize("mode", ["chat", "completion"]) +def test_non_stream_usage_keeps_existing_fallback_estimates( + http_server: HTTPServerFactory, + monkeypatch: pytest.MonkeyPatch, + mode: str, +) -> None: + choice = ( + {"message": {"content": "Answer"}} if mode == "chat" else {"text": "Answer"} + ) + server = http_server([Response(body=json.dumps({"choices": [choice]}).encode())]) + model = OAICompatLargeLanguageModel([]) + monkeypatch.setattr(model, "_get_num_tokens_by_gpt2", len) + + chunks = list( + model.invoke( + "model", + {"mode": mode, "endpoint_url": server.url}, + [ + UserPromptMessage(content="Question"), + AssistantPromptMessage(content="history"), + ], + stream=False, + ) + ) + + usage = chunks[-1].delta.usage + assert usage is not None + assert usage.prompt_tokens == 285 + assert usage.completion_tokens == 215 + assert model.started_at == 0 + + +@pytest.mark.parametrize("calling_type", ["tool_call", "function_call"]) +@pytest.mark.parametrize("reported_usage", [True, False]) +@pytest.mark.parametrize("include_content", [True, False]) +def test_non_stream_tool_calls_preserve_arguments_and_request_schema( + http_server: HTTPServerFactory, + monkeypatch: pytest.MonkeyPatch, + calling_type: str, + reported_usage: bool, + include_content: bool, +) -> None: + function = {"name": "weather", "arguments": '{"city":"北京"}'} + message = {"content": None} if include_content else {} + if calling_type == "tool_call": + message["tool_calls"] = [ + {"id": "call-1", "type": "function", "function": function} + ] + else: + message["function_call"] = function + response = {"choices": [{"message": message}]} + if reported_usage: + response["usage"] = {"prompt_tokens": 0, "completion_tokens": 0} + server = http_server([Response(body=json.dumps(response).encode())]) + tool = PromptMessageTool( + name="weather", description="Get weather", parameters={"type": "object"} + ) + + model = OAICompatLargeLanguageModel([]) + monkeypatch.setattr(model, "_get_num_tokens_by_gpt2", len) + invocation = model.invoke( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": calling_type, + }, + [UserPromptMessage(content="Weather?")], + tools=[tool], + stream=False, + ) + if include_content and reported_usage: + chunks = list(invocation) + call = chunks[-1].delta.message.tool_calls[0] + assert chunks[-1].delta.message.content is None + assert call.function.name == "weather" + assert json.loads(call.function.arguments) == {"city": "北京"} + usage = chunks[-1].delta.usage + assert usage is not None + assert (usage.prompt_tokens, usage.completion_tokens) == (0, 0) + else: + with pytest.raises(InvokeError) as exc_info: + list(invocation) + assert isinstance( + exc_info.value.__cause__, KeyError if not include_content else ValueError + ) + assert model.started_at == 0 + payload = json.loads(server.requests[0].body) + if calling_type == "tool_call": + assert payload["tools"] == [{"type": "function", "function": tool.model_dump()}] + else: + assert payload["functions"] == [tool.model_dump()] + + +@pytest.mark.parametrize( + ("delimiter", "configured"), [("\n\n", "\n\n"), ("\n", r"\n"), ("|", "|")] +) +def test_stream_decodes_fragmented_utf8_and_configured_delimiters( + http_server: HTTPServerFactory, delimiter: str, configured: str +) -> None: + events = [ + ": keepalive", + 'data: {"choices":[{"delta":{"reasoning_content":"想"}}]}', + 'data: {"choices":[{"delta":{}}]}', + 'data: {"choices":[{"delta":{"content":"回答😀"},"finish_reason":"stop"}]}', + 'data: {"choices":[],"usage":{"prompt_tokens":0,"completion_tokens":2}}', + "data: [DONE]", + "data: invalid-after-done", + ] + body = (delimiter.join(events) + delimiter).encode() + server = http_server([ + Response( + headers={"Content-Type": "text/event-stream"}, + chunks=tuple(body[index : index + 1] for index in range(len(body))), + ) + ]) + model = OAICompatLargeLanguageModel([]) + + chunks = list( + model.invoke( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "stream_mode_delimiter": configured, + }, + [UserPromptMessage(content="Question")], + ) + ) + + assert ( + "".join(chunk.delta.message.content or "" for chunk in chunks) + == "\n想\n回答😀" + ) + assert chunks[-1].delta.finish_reason == "stop" + assert chunks[-1].delta.usage is not None + assert chunks[-1].delta.usage.total_tokens == 2 + assert model.started_at == 0 + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("status", [401, 429, 503]) +def test_http_errors_reach_public_invoke_and_clear_timing( + http_server: HTTPServerFactory, stream: bool, status: int +) -> None: + server = http_server([Response(status=status, body=b"upstream rejected request")]) + model = OAICompatLargeLanguageModel([]) + + with pytest.raises(InvokeError, match=f"{status}: upstream rejected request"): + list( + model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + stream=stream, + ) + ) + + assert model.started_at == 0 + + +@pytest.mark.parametrize( + "error", [b"{broken", b'{"error":{"message":"upstream failed"}}'] +) +def test_stream_errors_after_content_do_not_report_success( + http_server: HTTPServerFactory, error: bytes +) -> None: + server = http_server([ + Response( + chunks=( + b'data: {"choices":[{"delta":{"content":"partial"}}]}\n\n', + b"data: " + error + b"\n\n", + ) + ) + ]) + model = OAICompatLargeLanguageModel([]) + stream = model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + + assert next(stream).delta.message.content == "partial" + with pytest.raises(InvokeError): + next(stream) + assert model.started_at == 0 + + +@pytest.mark.parametrize("ending", [b"", b"\n\n", b"\n\ndata: [DONE]\n\n"]) +def test_complete_http_response_keeps_eof_and_done_acceptance( + http_server: HTTPServerFactory, + ending: bytes, +) -> None: + server = http_server([ + Response( + chunks=( + ( + b'data: {"choices":[{"delta":{"content":"partial"},' + b'"finish_reason":null}]}' + ) + + ending, + ) + ) + ]) + model = OAICompatLargeLanguageModel([]) + stream = model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + + assert next(stream).delta.message.content == "partial" + terminal = next(stream) + assert terminal.delta.message == AssistantPromptMessage(content="") + assert terminal.delta.finish_reason is None + assert terminal.delta.usage is not None + assert list(stream) == [] + assert model.started_at == 0 + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("calling_type", ["tool_call", "function_call", "no_call"]) +def test_missing_usage_keeps_estimates_and_provider_response_overrides( + http_server: HTTPServerFactory, + monkeypatch: pytest.MonkeyPatch, + stream: bool, + calling_type: str, +) -> None: + response = ( + Response( + chunks=( + ( + b'data: {"choices":[{"delta":{"content":"Answer"},' + b'"finish_reason":"stop"}]}\n\n' + ), + ) + ) + if stream + else Response(body=b'{"choices":[{"message":{"content":"Answer"}}]}') + ) + server = http_server([response]) + model = LegacyResponseLLM([]) + monkeypatch.setattr(model, "_get_num_tokens_by_gpt2", len) + prompts = [UserPromptMessage(content="Question")] + tools = [ + PromptMessageTool( + name="search", + description="Search documentation. " * 100, + parameters={"type": "object"}, + ) + ] + credentials = { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": calling_type, + "input_price": "1", + "unit": "0.001", + } + expected = 18 if stream else 285 + + chunks = list( + model.invoke("model", credentials, prompts, tools=tools, stream=stream) + ) + + usage = chunks[-1].delta.usage + assert usage is not None + assert usage.prompt_tokens == expected + assert usage.prompt_price == Decimal(expected) * Decimal("0.001") + payload = json.loads(server.requests[0].body) + if calling_type == "tool_call": + assert payload["tools"][0]["function"]["description"] == tools[0].description + elif calling_type == "function_call": + assert payload["functions"][0]["description"] == tools[0].description + else: + assert "tools" not in payload + assert "functions" not in payload + + +@pytest.mark.parametrize( + ("mode", "stream"), [("chat", False), ("completion", False), ("chat", True)] +) +def test_validate_credentials_uses_configured_endpoint_and_mode( + http_server: HTTPServerFactory, mode: str, stream: bool +) -> None: + server = http_server([Response(body=b"unused stream body" if stream else b"{}")]) + OAICompatLargeLanguageModel([]).validate_credentials( + "model", + { + "mode": mode, + "endpoint_url": f"{server.url}/v1", + "stream_mode_auth": "use" if stream else "not_use", + "validate_credentials_max_tokens": 7, + "extra_headers": {"X-Dify-App-Id": "app-1"}, + }, + ) + + request = server.requests[0] + payload = json.loads(request.body) + assert request.path == ( + "/v1/chat/completions" if mode == "chat" else "/v1/completions" + ) + assert request.headers["X-Dify-App-Id"] == "app-1" + assert "Authorization" not in request.headers + assert payload["max_tokens"] == 7 + assert bool(payload.get("stream")) is stream + + +@pytest.mark.parametrize( + "response", + [ + Response(status=401, body=b"invalid credentials"), + Response(body=b"not json"), + Response(body=b'{"object":"embedding"}'), + Response(body=b"[]"), + ], +) +def test_validate_credentials_rejects_http_and_response_errors( + http_server: HTTPServerFactory, response: Response +) -> None: + server = http_server([response]) + with pytest.raises(CredentialsValidateFailedError): + OAICompatLargeLanguageModel([]).validate_credentials( + "model", {"mode": "chat", "endpoint_url": server.url} + ) + + +def test_stream_early_close_releases_the_http_response( + http_server: HTTPServerFactory, http_responses: tuple[list[HTTPResponse], Event] +) -> None: + responses, _ready = http_responses + server = http_server([ + Response( + chunks=( + b'data: {"choices":[{"delta":{"content":"partial"}}]}\n\n', + b"data: [DONE]\n\n", + ) + ) + ]) + model = OAICompatLargeLanguageModel([]) + stream = model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + + assert next(stream).delta.message.content == "partial" + stream.close() + + assert responses[0].raw.closed + assert model.started_at == 0 + + +def test_stream_greenlet_cancellation_releases_the_http_response( + http_server: HTTPServerFactory, http_responses: tuple[list[HTTPResponse], Event] +) -> None: + responses, ready = http_responses + server = http_server([Response(chunks=(b"data: [DONE]\n\n",), delay=10)]) + model = OAICompatLargeLanguageModel([]) + worker = gevent.spawn( + list, + model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ), + ) + try: + assert ready.wait(timeout=2), "The HTTP response was not received" + finally: + worker.kill(timeout=2) + + assert worker.ready() + assert responses[0].raw.closed + assert model.started_at == 0 + + +def test_stream_read_timeout_closes_the_response_and_clears_timing( + http_server: HTTPServerFactory, + http_responses: tuple[list[HTTPResponse], Event], + monkeypatch: pytest.MonkeyPatch, +) -> None: + responses, _ready = http_responses + server = http_server([Response(chunks=(b"data: [DONE]\n\n",), delay=10)]) + monkeypatch.setattr( + llm, + "_plugin_config", + llm._plugin_config.model_copy(update={"MAX_REQUEST_TIMEOUT": 1}), + ) + model = OAICompatLargeLanguageModel([]) + + with pytest.raises(InvokeError, match="timed out"): + list( + model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + ) + + assert responses[0].raw.closed + assert model.started_at == 0 + + +@pytest.mark.parametrize( + "body", + [ + b"not json", + b"{}", + b"[]", + b'{"choices":[]}', + b'{"choices":[{}],"usage":{"prompt_tokens":11,"completion_tokens":0}}', + ], +) +def test_non_stream_malformed_response_fails_without_leaking_timing( + http_server: HTTPServerFactory, body: bytes +) -> None: + server = http_server([Response(body=body)]) + model = OAICompatLargeLanguageModel([]) + with pytest.raises(InvokeError): + list( + model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + stream=False, + ) + ) + assert model.started_at == 0 + + +def test_redirect_preserves_the_request_body(http_server: HTTPServerFactory) -> None: + server = http_server([ + Response(status=307, headers={"Location": "/v2/chat/completions"}), + Response( + body=b'{"choices":[{"message":{"content":"answer"}}],"usage":{"prompt_tokens":1,"completion_tokens":1}}' + ), + ]) + + chunks = list( + OAICompatLargeLanguageModel([]).invoke( + "model", + {"mode": "chat", "endpoint_url": server.url, "api_key": "key"}, + [UserPromptMessage(content="Question")], + stream=False, + ) + ) + + assert chunks[0].delta.message.content == "answer" + first, redirected = server.requests + assert redirected.path == "/v2/chat/completions" + assert redirected.method == "POST" + assert redirected.body == first.body + assert redirected.headers["Authorization"] == "Bearer key" + + +@pytest.mark.parametrize( + ("newline", "prefix"), + [ + ("\n", ""), + ("\r\n", ""), + ("\r", ""), + ("\n", ": keepalive\n"), + ("\n", "event: message\n"), + ], + ids=["lf", "crlf", "cr", "legacy-leading-comment", "event-type"], +) +def test_stream_handles_sse_framing_and_legacy_leading_comments( + http_server: HTTPServerFactory, + http_responses: tuple[list[HTTPResponse], Event], + newline: str, + prefix: str, +) -> None: + responses, _ready = http_responses + events = [ + prefix + + 'data: {"choices":[{"delta":{"content":"Answer"},"finish_reason":"stop"}]}', + 'data: {"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}', + "data: [DONE]", + "", + ] + body = (newline * 2).join(events).encode() + server = http_server([ + Response( + headers={"Content-Type": "text/event-stream"}, + chunks=tuple( + body[offset : offset + 3] for offset in range(0, len(body), 3) + ), + ) + ]) + model = OAICompatLargeLanguageModel([]) + + try: + chunks = list( + model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + ) + finally: + assert responses[0].raw.closed + assert model.started_at == 0 + + # Preserve the baseline's discard for leading comments until #397 is addressed. + assert "".join(chunk.delta.message.get_text_content() for chunk in chunks) == ( + "" if prefix.startswith(":") else "Answer" + ) + assert chunks[-1].delta.finish_reason == ( + None if prefix.startswith(":") else "stop" + ) + assert chunks[-1].delta.usage is not None + assert chunks[-1].delta.usage.total_tokens == 3 + + +@pytest.mark.parametrize("charset", ["utf-8", "gbk", "utf-16"]) +def test_stream_preserves_explicit_response_charset( + http_server: HTTPServerFactory, charset: str +) -> None: + body = ( + 'data: {"choices":[{"delta":{"content":"回答"},"finish_reason":"stop"}]}' + '\n\ndata: {"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}' + "\n\ndata: [DONE]\n\n" + ).encode(charset) + server = http_server([ + Response( + headers={"Content-Type": f"text/event-stream; charset={charset}"}, + chunks=tuple(body[index : index + 1] for index in range(len(body))), + ) + ]) + model = OAICompatLargeLanguageModel([]) + + chunks = list( + model.invoke( + "model", + {"mode": "chat", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + ) + + assert "".join(chunk.delta.message.get_text_content() for chunk in chunks) == "回答" + assert chunks[-1].delta.usage is not None + assert chunks[-1].delta.usage.total_tokens == 3 + assert model.started_at == 0 + + +@pytest.mark.parametrize("delimiter", ["\n\n", "|"]) +@pytest.mark.parametrize("prefix", ["", "data: "]) +def test_stream_preserves_raw_and_single_prefix_multiline_json( + http_server: HTTPServerFactory, delimiter: str, prefix: str +) -> None: + payload = { + "choices": [{"delta": {"content": "Answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2}, + } + body = prefix + json.dumps(payload, indent=2) + delimiter + "[DONE]" + delimiter + server = http_server([Response(body=body.encode())]) + model = OAICompatLargeLanguageModel([]) + + chunks = list( + model.invoke( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "stream_mode_delimiter": delimiter, + }, + [UserPromptMessage(content="Question")], + ) + ) + + assert ( + "".join(chunk.delta.message.get_text_content() for chunk in chunks) == "Answer" + ) + assert chunks[-1].delta.finish_reason == "stop" + assert chunks[-1].delta.usage is not None + assert chunks[-1].delta.usage.total_tokens == 3 + assert model.started_at == 0 + + +@pytest.mark.parametrize("delimiter", ["\n\n", "|"]) +@pytest.mark.parametrize( + "raw", + ["not json", "0", "false", '"text"', "[1]", 'not json\ndata: {"choices":[]}'], +) +def test_stream_keeps_errors_for_raw_nonobject_frames( + http_server: HTTPServerFactory, + http_responses: tuple[list[HTTPResponse], Event], + delimiter: str, + raw: str, +) -> None: + responses, _ready = http_responses + body = ( + 'data: {"choices":[{"delta":{"content":"partial"}}]}' + + delimiter + + raw + + delimiter + ).encode() + server = http_server([Response(body=body)]) + model = OAICompatLargeLanguageModel([]) + stream = model.invoke( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "stream_mode_delimiter": delimiter, + }, + [UserPromptMessage(content="Question")], + ) + + assert next(stream).delta.message.content == "partial" + with pytest.raises(InvokeError, match="Invalid JSON"): + next(stream) + assert responses[0].raw.closed + assert model.started_at == 0 + + +def test_completion_stream_handles_empty_text_and_partial_usage( + http_server: HTTPServerFactory, +) -> None: + frames = [ + {"choices": [{"text": ""}]}, + {"choices": [{"text": "An"}]}, + {"choices": [{"text": ""}]}, + {"choices": [{"text": "swer"}]}, + { + "choices": [{"text": "", "finish_reason": "length"}], + "usage": {"prompt_tokens": 0}, + }, + ] + body = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + server = http_server([Response(body=(body + "data: [DONE]\n\n").encode())]) + model = OAICompatLargeLanguageModel([]) + + chunks = list( + model.invoke( + "model", + {"mode": "completion", "endpoint_url": server.url}, + [UserPromptMessage(content="Question")], + ) + ) + + assert ( + "".join(chunk.delta.message.get_text_content() for chunk in chunks) == "Answer" + ) + assert chunks[-1].delta.finish_reason == "length" + usage = chunks[-1].delta.usage + assert usage is not None + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + 0, + 1, + 1, + ) + assert model.started_at == 0 + + +@pytest.mark.parametrize("stream", [False, True]) +def test_missing_usage_can_estimate_tools_with_integer_enum_parameters( + http_server: HTTPServerFactory, stream: bool +) -> None: + response = ( + Response( + body=b'data: {"choices":[{"delta":{"content":"Answer"},' + b'"finish_reason":"stop"}]}\n\n' + b"data: [DONE]\n\n" + ) + if stream + else Response(body=b'{"choices":[{"message":{"content":"Answer"}}]}') + ) + server = http_server([response]) + model = OAICompatLargeLanguageModel([]) + tool = PromptMessageTool( + name="configure", + description="Choose the number of attempts", + parameters={ + "type": "object", + "properties": {"attempts": {"type": "integer", "enum": [1, 2]}}, + "required": ["attempts"], + }, + ) + + chunks = list( + model.invoke( + "model", + { + "mode": "chat", + "endpoint_url": server.url, + "function_calling_type": "tool_call", + }, + [UserPromptMessage(content="Configure two attempts")], + tools=[tool], + stream=stream, + ) + ) + + assert ( + "".join(chunk.delta.message.get_text_content() for chunk in chunks) == "Answer" + ) + assert chunks[-1].delta.usage is not None + assert chunks[-1].delta.usage.prompt_tokens > 0 + payload = json.loads(server.requests[0].body) + assert payload["tools"] == [{"type": "function", "function": tool.model_dump()}] + assert model.started_at == 0 diff --git a/tests/interfaces/model/openai_compatible/test_llm.py b/tests/interfaces/model/openai_compatible/test_llm.py index 688ef92f..41edbe4a 100644 --- a/tests/interfaces/model/openai_compatible/test_llm.py +++ b/tests/interfaces/model/openai_compatible/test_llm.py @@ -1,15 +1,14 @@ import json -from collections import UserDict -from collections.abc import Mapping from http import HTTPStatus +from io import BytesIO from types import SimpleNamespace -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import MagicMock, patch import gevent.socket import pytest +import requests from gevent.threadpool import ThreadPool -from dify_plugin.entities.model.llm import LLMResultChunk from dify_plugin.entities.model.message import ( AssistantPromptMessage, AudioPromptMessageContent, @@ -22,42 +21,13 @@ UserPromptMessage, VideoPromptMessageContent, ) -from dify_plugin.errors.model import CredentialsValidateFailedError +from dify_plugin.errors.model import CredentialsValidateFailedError, InvokeError from dify_plugin.interfaces.model import ai_model from dify_plugin.interfaces.model.openai_compatible.llm import ( OAICompatLargeLanguageModel, ) -class DuckJsonObject: - def __init__(self) -> None: - self.trace: list[str] = [] - - def get(self, key: str, default: object = None) -> str: - del key, default - self.trace.append("get") - return "chat.completion" - - def __contains__(self, key: object) -> bool: - del key - self.trace.append("contains") - return True - - def __getitem__(self, key: str) -> str: - del key - self.trace.append("getitem") - return "chat.completion" - - -class StatefulTruth: - def __init__(self) -> None: - self.calls = 0 - - def __bool__(self) -> bool: - self.calls += 1 - return self.calls == 1 - - def _stream_choice(delta: dict, finish_reason: str | None = None) -> str: return "data: " + json.dumps({ "choices": [{"delta": delta, "finish_reason": finish_reason}] @@ -282,7 +252,7 @@ def test_convert_prompt_messages_preserves_text_and_empty_content() -> None: "content_type", [AudioPromptMessageContent, DocumentPromptMessageContent] ) @pytest.mark.parametrize("mixed", [False, True]) -def test_generate_preserves_unsupported_non_user_content_until_json_encoding( +def test_unsupported_non_user_content_is_rejected_before_sending_request( message_type: type[PromptMessage], content_type: type[AudioPromptMessageContent | DocumentPromptMessageContent], mixed: bool, @@ -313,105 +283,39 @@ def test_generate_preserves_unsupported_non_user_content_until_json_encoding( "function_calling_type": "tool_call", } llm = OAICompatLargeLanguageModel([]) - converted = llm._convert_prompt_message_to_dict(message, credentials) - assert converted["content"][-1] is unsupported with ( patch( "dify_plugin.interfaces.model.openai_compatible.llm.requests.post" ) as post, - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.json.dumps", - wraps=json.dumps, - ) as dumps, - pytest.raises( - TypeError, match=f"{content_type.__name__} is not JSON serializable" - ), + pytest.raises(InvokeError), ): - llm._generate("model", credentials, [message], {}, stream=False) + list(llm.invoke("model", credentials, [message], {}, stream=False)) - dumps.assert_called_once() post.assert_not_called() - assert message.content == content - - -@pytest.mark.parametrize( - "content_type", [AudioPromptMessageContent, DocumentPromptMessageContent] -) -def test_provider_override_can_serialize_unsupported_content_after_super( - content_type: type[AudioPromptMessageContent | DocumentPromptMessageContent], -) -> None: - class MediaUrlModel(OAICompatLargeLanguageModel): - def _convert_prompt_message_to_dict( - self, message: PromptMessage, credentials: dict | None = None - ) -> dict: - converted = super()._convert_prompt_message_to_dict(message, credentials) - converted["content"] = converted["content"][0].data - return converted - - content = content_type( - format="bin", - mime_type="application/octet-stream", - url="https://example.com/file", - ) - llm = MediaUrlModel([]) - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=MagicMock(status_code=HTTPStatus.OK), - ) as post, - patch.object(llm, "_handle_generate_response"), - ): - llm._generate( - "model", - {"endpoint_url": "https://example.com/v1", "mode": "chat"}, - [SystemPromptMessage(content=[content])], - {}, - stream=False, - ) - - assert json.loads(post.call_args.kwargs["data"])["messages"] == [ - {"role": "system", "content": content.data} - ] @pytest.mark.parametrize( "message_type", [SystemPromptMessage, AssistantPromptMessage, ToolPromptMessage] ) -def test_num_tokens_counts_non_user_text_without_counting_media( +def test_num_tokens_counts_non_user_text_parts_like_plain_text( message_type: type[PromptMessage], ) -> None: message = message_type.model_validate({ "content": [ TextPromptMessageContent(data="first"), - ImagePromptMessageContent( - format="png", mime_type="image/png", url="https://example.com/image.png" - ), TextPromptMessageContent(data="second"), - AudioPromptMessageContent( - format="mp3", - mime_type="audio/mpeg", - url="https://example.com/audio.mp3", - ), - DocumentPromptMessageContent( - format="pdf", - mime_type="application/pdf", - url="https://example.com/doc.pdf", - ), ], "tool_call_id": "call-1", }) llm = OAICompatLargeLanguageModel([]) - with patch.object(llm, "_get_num_tokens_by_gpt2", side_effect=len) as tokenize: - count = llm._num_tokens_from_messages( - [message], credentials={"function_calling_type": "tool_call"} - ) + credentials = {"function_calling_type": "tool_call"} + count = llm.get_num_tokens("model", credentials, [message]) + text_only = message.model_copy(update={"content": "firstsecond"}) + empty = message.model_copy(update={"content": ""}) - expected = [message.role.value, "firstsecond"] - if isinstance(message, ToolPromptMessage): - expected.append("call-1") - assert [call.args[0] for call in tokenize.call_args_list] == expected - assert count == 6 + sum(map(len, expected)) + assert count == llm.get_num_tokens("model", credentials, [text_only]) + assert count > llm.get_num_tokens("model", credentials, [empty]) @pytest.mark.parametrize("use_threadpool", [False, True]) @@ -432,15 +336,14 @@ def test_num_tokens_treats_special_token_spelling_as_text( ToolPromptMessage(content=content, tool_call_id="call-1"), ] try: - count = OAICompatLargeLanguageModel([]).get_num_tokens( - "model", {"function_calling_type": "tool_call"}, messages - ) + model = OAICompatLargeLanguageModel([]) + credentials = {"function_calling_type": "tool_call"} + count = model.get_num_tokens("model", credentials, messages) + empty = [message.model_copy(update={"content": ""}) for message in messages] + assert count > model.get_num_tokens("model", credentials, empty) finally: pool.kill() - # GPT2 counts the literal spelling as ordinary text, not a control token. - assert count == 55 - def test_generate_encodes_request_json_as_utf8() -> None: content = "你好😀\ud800" @@ -510,229 +413,6 @@ def test_validate_credentials_rejects_non_mapping_extra_headers( post.assert_not_called() -def test_validate_credentials_wraps_unreadable_error_response() -> None: - error = RuntimeError("broken body") - response = MagicMock(status_code=HTTPStatus.BAD_REQUEST) - type(response).text = PropertyMock(side_effect=error) - response.close.side_effect = RuntimeError("close failed") - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - "stream_mode_auth": "use", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(RuntimeError) as exc_info, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert exc_info.value is error - response.close.assert_called_once() - - -def test_validate_credentials_accepts_mapping_response() -> None: - response = MagicMock(status_code=HTTPStatus.OK) - response.json.return_value = UserDict({"object": "chat.completion"}) - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - -def test_validate_credentials_wraps_lazy_mapping_errors() -> None: - error_message = "lazy response failed" - json_result = MagicMock(spec=Mapping) - json_result.get.side_effect = RuntimeError(error_message) - response = MagicMock(status_code=HTTPStatus.OK) - response.json.return_value = json_result - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError, match=error_message), - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - -def test_validate_credentials_preserves_prepared_credential_errors() -> None: - sentinel = CredentialsValidateFailedError("prepared credentials failed") - credentials = MagicMock() - credentials.get.side_effect = sentinel - - with pytest.raises(CredentialsValidateFailedError) as exc_info: - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert exc_info.value is sentinel - - -def test_validate_credentials_preserves_response_credential_errors() -> None: - sentinel = CredentialsValidateFailedError("response failed") - response = MagicMock(status_code=HTTPStatus.OK) - response.json.side_effect = sentinel - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError) as exc_info, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert exc_info.value is sentinel - - -def test_validate_credentials_preserves_mapping_membership_access() -> None: - json_result = DuckJsonObject() - response = MagicMock(status_code=HTTPStatus.OK) - response.json.return_value = json_result - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert json_result.trace == ["get", "contains", "getitem"] - - -def test_validate_credentials_compares_status_once() -> None: - status = MagicMock() - status.__ne__.side_effect = [True, False] - response = MagicMock(status_code=status, text="failed") - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - "stream_mode_auth": "use", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError), - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - status.__ne__.assert_called_once_with(HTTPStatus.OK) - - -def test_validate_credentials_preserves_failed_status_second_access() -> None: - sentinel = CredentialsValidateFailedError("second status read failed") - response = MagicMock(text="failed") - type(response).status_code = PropertyMock( - side_effect=[HTTPStatus.BAD_REQUEST, sentinel], - ) - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError) as exc_info, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert exc_info.value is sentinel - - -def test_validate_credentials_preserves_response_truthiness_hook() -> None: - sentinel = CredentialsValidateFailedError("response truthiness failed") - response = MagicMock(status_code=HTTPStatus.OK) - response.json.side_effect = RuntimeError("json failed") - response.__bool__.side_effect = sentinel - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError) as exc_info, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert exc_info.value is sentinel - - -def test_validate_credentials_preserves_status_failure_truthiness_hook() -> None: - sentinel = CredentialsValidateFailedError("response truthiness failed") - response = MagicMock(text="failed") - type(response).status_code = PropertyMock( - side_effect=[HTTPStatus.BAD_REQUEST, RuntimeError("status failed")], - ) - response.__bool__.side_effect = sentinel - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - } - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError) as exc_info, - ): - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert exc_info.value is sentinel - - -def test_validate_credentials_truth_tests_stream_mode_once() -> None: - stream_result = StatefulTruth() - stream_mode = MagicMock() - stream_mode.__eq__.return_value = stream_result - response = MagicMock(status_code=HTTPStatus.OK) - credentials = { - "endpoint_url": "https://example.com/v1", - "mode": "chat", - "stream_mode_auth": stream_mode, - } - - with patch( - "dify_plugin.interfaces.model.openai_compatible.llm.requests.post", - return_value=response, - ) as post: - OAICompatLargeLanguageModel([]).validate_credentials("model", credentials) - - assert stream_result.calls == 1 - assert post.call_args.kwargs["stream"] is True - response.json.assert_not_called() - - @pytest.mark.parametrize( ("lines", "expected_content", "expected_finish_reason", "expected_usage"), [ @@ -858,35 +538,30 @@ def test_stream_reasoning_is_closed_at_end( expected_finish_reason: str | None, expected_usage: dict | None, ) -> None: - response = MagicMock() - response.iter_lines.return_value = lines + response = requests.Response() + response.encoding = "utf-8" + response.raw = BytesIO(("\n\n".join(lines) + "\n\n").encode()) llm = OAICompatLargeLanguageModel([]) - final_chunk = object() - - with patch.object( - llm, - "_create_final_llm_result_chunk", - return_value=final_chunk, - ) as create_final: - results = list( - llm._handle_generate_stream_response( - model="model", - credentials={}, - response=response, - prompt_messages=[], - ) + results = list( + llm._handle_generate_stream_response( + model="model", + credentials={}, + response=response, + prompt_messages=[], ) + ) - stream_chunks = [result for result in results if isinstance(result, LLMResultChunk)] - assert "".join(chunk.delta.message.content for chunk in stream_chunks) == ( + assert "".join(chunk.delta.message.content or "" for chunk in results) == ( expected_content ) - assert [chunk.delta.index for chunk in stream_chunks] == sorted( - chunk.delta.index for chunk in stream_chunks + assert [chunk.delta.index for chunk in results] == sorted( + chunk.delta.index for chunk in results ) - assert results[-1] is final_chunk - create_final.assert_called_once() - assert create_final.call_args.kwargs["full_content"] == expected_content - assert create_final.call_args.kwargs["finish_reason"] == expected_finish_reason - assert create_final.call_args.kwargs["usage"] == expected_usage - assert create_final.call_args.kwargs["index"] > stream_chunks[-1].delta.index + terminal = results[-1].delta + assert terminal.finish_reason == expected_finish_reason + assert terminal.usage is not None + if expected_usage is not None: + assert terminal.usage.prompt_tokens == expected_usage["prompt_tokens"] + assert terminal.usage.completion_tokens == expected_usage["completion_tokens"] + else: + assert terminal.usage.completion_tokens > 0 diff --git a/tests/interfaces/model/openai_compatible/test_llm_stream_lifecycle.py b/tests/interfaces/model/openai_compatible/test_llm_stream_lifecycle.py index a76e8aea..d29942f0 100644 --- a/tests/interfaces/model/openai_compatible/test_llm_stream_lifecycle.py +++ b/tests/interfaces/model/openai_compatible/test_llm_stream_lifecycle.py @@ -4,6 +4,8 @@ import pytest import requests +from hypothesis import given +from hypothesis import strategies as st from dify_plugin.entities.model.llm import LLMResultChunk from dify_plugin.entities.model.message import UserPromptMessage @@ -275,3 +277,45 @@ def test_stream_merges_partial_usage_without_losing_counts( assert usage.total_tokens == prompt_tokens assert chunks[-1].delta.finish_reason == "stop" response.close.assert_called_once() + + +@given( + content=st.text(alphabet=st.characters(exclude_categories=("Cs",)), max_size=100), + sizes=st.lists(st.integers(min_value=1, max_value=32), min_size=1, max_size=20), + newline=st.sampled_from(["\n", "\r\n", "\r"]), +) +def test_stream_decoding_is_independent_of_byte_boundaries( + content: str, + sizes: list[int], + newline: str, +) -> None: + payload = {"choices": [{"delta": {"content": content}, "finish_reason": "stop"}]} + data_lines = "\n".join( + f"data: {line}" + for line in json.dumps(payload, ensure_ascii=False, indent=2).split("\n") + ) + body = ( + ( + f"\ufeffevent: message\nid: 1\n: heartbeat\n{data_lines}\n\n" + + ZERO_USAGE.decode() + + "data: [DONE]\n\n" + ) + .replace("\n", newline) + .encode() + ) + fragments = [] + offset = 0 + for size in sizes: + fragments.append(body[offset : offset + size]) + offset += size + fragments.append(body[offset:]) + response = _response(fragments) + model = OAICompatLargeLanguageModel([]) + + with model.timing_context(): + chunks = list(_stream(model, response)) + + assert "".join(chunk.delta.message.content for chunk in chunks) == content + assert chunks[-1].delta.finish_reason == "stop" + assert chunks[-1].delta.usage.total_tokens == 0 + response.close.assert_called_once() diff --git a/tests/interfaces/model/openai_compatible/test_model_configuration.py b/tests/interfaces/model/openai_compatible/test_model_configuration.py new file mode 100644 index 00000000..dcd610ef --- /dev/null +++ b/tests/interfaces/model/openai_compatible/test_model_configuration.py @@ -0,0 +1,192 @@ +from decimal import Decimal + +import pytest + +from dify_plugin.entities.model import ModelFeature, ModelPropertyKey +from dify_plugin.entities.model.message import ( + AssistantPromptMessage, + ImagePromptMessageContent, + PromptMessageTool, + TextPromptMessageContent, + ToolPromptMessage, + UserPromptMessage, +) +from dify_plugin.interfaces.model.openai_compatible.llm import ( + OAICompatLargeLanguageModel, +) +from dify_plugin.interfaces.model.openai_compatible.text_embedding import ( + OAICompatEmbeddingModel, +) + + +@pytest.mark.parametrize("calling_type", ["tool_call", "function_call"]) +def test_token_budget_includes_tool_history_and_tool_definitions( + calling_type: str, +) -> None: + model = OAICompatLargeLanguageModel([]) + credentials = {"function_calling_type": calling_type} + prompt = UserPromptMessage(content="What is the weather in Paris?") + history = [ + prompt, + AssistantPromptMessage( + tool_calls=[ + AssistantPromptMessage.ToolCall( + id="weather", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction( + name="weather", arguments='{"city":"Paris"}' + ), + ) + ] + ), + ToolPromptMessage(content="It is sunny.", tool_call_id="weather"), + ] + minimal_tool = PromptMessageTool( + name="weather", description="", parameters={"type": "object"} + ) + tool = PromptMessageTool( + name="weather", + description="Get the current weather for a city.", + parameters={ + "title": "Weather input", + "type": "object", + "properties": { + "city": {"type": "string", "description": "The city to look up."}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["city"], + }, + ) + + prompt_tokens = model.get_num_tokens("model", credentials, [prompt]) + history_tokens = model.get_num_tokens("model", credentials, history) + minimal_tool_tokens = model.get_num_tokens( + "model", credentials, history, [minimal_tool] + ) + full_tool_tokens = model.get_num_tokens("model", credentials, history, [tool]) + + assert prompt_tokens < history_tokens < minimal_tool_tokens < full_tool_tokens + assert ( + model.get_num_tokens("model", credentials, history, [tool, minimal_tool]) + > full_tool_tokens + ) + + +def test_token_budget_counts_names_and_text_without_encoding_image_urls() -> None: + model = OAICompatLargeLanguageModel([]) + prompt = UserPromptMessage(content="Describe this image") + image = ImagePromptMessageContent( + format="png", mime_type="image/png", url="https://example.com/image.png" + ) + image_prompt = UserPromptMessage( + content=[TextPromptMessageContent(data="Describe this image"), image] + ) + text_tokens = model.get_num_tokens("model", {}, [prompt]) + image_tokens = model.get_num_tokens("model", {}, [image_prompt]) + long_url_prompt = image_prompt.model_copy(deep=True) + long_url_prompt.content[1].url += "?signature=" + "a" * 10_000 + + assert model.get_num_tokens("model", {}, [long_url_prompt]) == image_tokens + assert image_tokens >= text_tokens + prompt.name = "planner" + assert model.get_num_tokens("model", {}, [prompt]) > text_tokens + + +@pytest.mark.parametrize( + ("calling_type", "media", "stream_tools", "features"), + [ + ( + "function_call", + "support", + "supported", + { + ModelFeature.TOOL_CALL, + ModelFeature.STREAM_TOOL_CALL, + ModelFeature.VISION, + ModelFeature.VIDEO, + }, + ), + ("tool_call", "not_support", "not_supported", {ModelFeature.MULTI_TOOL_CALL}), + ("no_call", "not_support", "not_supported", set()), + ], +) +def test_custom_model_capabilities_and_limits_follow_credentials( + calling_type: str, media: str, stream_tools: str, features: set[ModelFeature] +) -> None: + schema = OAICompatLargeLanguageModel([]).get_model_schema( + "custom-model", + { + "mode": "chat", + "function_calling_type": calling_type, + "vision_support": media, + "video_support": media, + "stream_function_calling": stream_tools, + "context_size": "32768", + "max_tokens_to_sample": "2048", + "input_price": "2.5", + "output_price": "5", + "unit": "0.001", + }, + ) + assert schema is not None + assert set(schema.features or []) == features + assert schema.model_properties[ModelPropertyKey.CONTEXT_SIZE] == 32768 + max_tokens = next( + rule for rule in schema.parameter_rules if rule.name == "max_tokens" + ) + assert max_tokens.max == 2048 + assert schema.pricing is not None + assert schema.pricing.input == Decimal("2.5") + assert schema.pricing.output == Decimal(5) + + +def test_embedding_token_estimates_preserve_input_order_and_empty_texts() -> None: + model = OAICompatEmbeddingModel([]) + assert model.get_num_tokens("model", {}, ["Hello", "", "Hello world"]) == [1, 0, 2] + + +def test_large_prompts_have_positive_nondecreasing_token_estimates() -> None: + model = OAICompatLargeLanguageModel([]) + + counts = [ + model.get_num_tokens("model", {}, [UserPromptMessage(content="a" * size)]) + for size in (100_000, 200_000) + ] + + assert all(type(count) is int and count > 0 for count in counts) + assert counts[1] >= counts[0] + + +@pytest.mark.parametrize( + "parameter", + [ + {"type": "string", "enum": ["fast", "accurate"]}, + {"type": "integer", "enum": [1, 2]}, + {"type": "boolean", "enum": [True, False]}, + {"type": ["string", "null"], "enum": ["automatic", None]}, + ], + ids=["string", "integer", "boolean", "nullable"], +) +def test_token_budget_accepts_json_schema_enum_values(parameter: dict) -> None: + model = OAICompatLargeLanguageModel([]) + prompts = [UserPromptMessage(content="Choose a setting")] + tool = PromptMessageTool( + name="configure", + description="Configure the setting", + parameters={ + "type": "object", + "properties": {"setting": parameter}, + "required": ["setting"], + }, + ) + + assert model.get_num_tokens("model", {}, prompts, [tool]) > model.get_num_tokens( + "model", {}, prompts + ) + + +def test_invalid_custom_model_mode_reports_a_configuration_error() -> None: + with pytest.raises(ValueError, match="invalid-mode"): + OAICompatLargeLanguageModel([]).get_model_schema( + "model", {"mode": "invalid-mode"} + ) diff --git a/tests/interfaces/model/openai_compatible/test_plugin_config.py b/tests/interfaces/model/openai_compatible/test_plugin_config.py deleted file mode 100644 index d7958e2a..00000000 --- a/tests/interfaces/model/openai_compatible/test_plugin_config.py +++ /dev/null @@ -1,75 +0,0 @@ -""" -Tests for _plugin_config configuration object in openai_compatible/llm.py - -These tests verify: -1. Configuration object is properly instantiated at module level -2. Configuration object is a singleton (only one instance at module level) -3. Configuration object has correct default values -4. Configuration object has correct types -""" - -import importlib - -from dify_plugin.config.config import DifyPluginEnv -from dify_plugin.interfaces.model.openai_compatible import llm - - -def test_plugin_config_exists() -> None: - """Test that _plugin_config object exists""" - assert hasattr(llm, "_plugin_config") - assert llm._plugin_config is not None - - -def test_plugin_config_type() -> None: - """Test that _plugin_config is an instance of DifyPluginEnv""" - assert isinstance(llm._plugin_config, DifyPluginEnv) - - -def test_plugin_config_has_max_request_timeout() -> None: - """Test that _plugin_config has MAX_REQUEST_TIMEOUT attribute""" - assert hasattr(llm._plugin_config, "MAX_REQUEST_TIMEOUT") - assert isinstance(llm._plugin_config.MAX_REQUEST_TIMEOUT, int) - # Default value should be 300 seconds (unless overridden by environment variable) - assert llm._plugin_config.MAX_REQUEST_TIMEOUT > 0 - - -def test_plugin_config_has_max_invocation_timeout() -> None: - """Test that _plugin_config has MAX_INVOCATION_TIMEOUT attribute""" - assert hasattr(llm._plugin_config, "MAX_INVOCATION_TIMEOUT") - assert isinstance(llm._plugin_config.MAX_INVOCATION_TIMEOUT, int) - # Default value should be 250 seconds (unless overridden by environment variable) - assert llm._plugin_config.MAX_INVOCATION_TIMEOUT > 0 - - -def test_plugin_config_singleton() -> None: - """Test that configuration object is instantiated only once at module level""" - # Get the id of the config object - config_id_1 = id(llm._plugin_config) - - # Import again, should be the same object - llm2 = importlib.import_module("dify_plugin.interfaces.model.openai_compatible.llm") - - config_id_2 = id(llm2._plugin_config) - - # Verify it's the same object - assert config_id_1 == config_id_2 - - -def test_plugin_config_timeout_values_reasonable() -> None: - """Test that timeout configuration values are within reasonable range""" - # MAX_REQUEST_TIMEOUT should be between 1 and 3600 seconds - assert 1 <= llm._plugin_config.MAX_REQUEST_TIMEOUT <= 3600 - - # MAX_INVOCATION_TIMEOUT should be between 1 and 3600 seconds - assert 1 <= llm._plugin_config.MAX_INVOCATION_TIMEOUT <= 3600 - - -def test_plugin_config_can_create_new_instance() -> None: - """Test that new DifyPluginEnv instances can be created. - - Verify the config class itself works. - """ - new_config = DifyPluginEnv() - assert isinstance(new_config, DifyPluginEnv) - assert hasattr(new_config, "MAX_REQUEST_TIMEOUT") - assert hasattr(new_config, "MAX_INVOCATION_TIMEOUT") diff --git a/tests/interfaces/model/openai_compatible/test_text_embedding.py b/tests/interfaces/model/openai_compatible/test_text_embedding.py index 7d9fa85b..e0a70701 100644 --- a/tests/interfaces/model/openai_compatible/test_text_embedding.py +++ b/tests/interfaces/model/openai_compatible/test_text_embedding.py @@ -1,7 +1,7 @@ -from http import HTTPStatus -from unittest.mock import MagicMock, PropertyMock, patch +from unittest.mock import Mock import pytest +import requests from dify_plugin.errors.model import CredentialsValidateFailedError from dify_plugin.interfaces.model.openai_compatible.text_embedding import ( @@ -9,131 +9,33 @@ ) -class DuckEmbeddingResponse: - def __init__(self) -> None: - self.trace: list[str] = [] +@pytest.mark.parametrize("error", [requests.ConnectionError, requests.Timeout]) +def test_credential_request_failure_keeps_the_original_cause( + monkeypatch: pytest.MonkeyPatch, error: type[requests.RequestException] +) -> None: + failure = error("provider unavailable") + monkeypatch.setattr(requests, "post", Mock(side_effect=failure)) - def __contains__(self, key: object) -> bool: - del key - self.trace.append("contains") - return True - - -class UnstringableResponseError(Exception): - def __init__(self, error: RuntimeError) -> None: - super().__init__() - self.error = error - - def __str__(self) -> str: - raise self.error - - -def test_validate_credentials_rejects_non_object_response() -> None: - response = MagicMock(status_code=HTTPStatus.OK) - response.json.return_value = 1 - response.close.side_effect = OSError("close failed") - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.text_embedding.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError), - ): + with pytest.raises(CredentialsValidateFailedError) as caught: OAICompatEmbeddingModel([]).validate_credentials( "model", {"endpoint_url": "https://example.com/v1"} ) - response.close.assert_called_once_with() + assert caught.value.__cause__ is failure -def test_validate_credentials_wraps_response_inspection_errors() -> None: - response = MagicMock() - type(response).status_code = PropertyMock( - side_effect=RuntimeError("broken status"), - ) +def test_response_close_failure_keeps_the_validation_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response = requests.Response() + response.status_code = 401 + close = Mock(side_effect=OSError("close failed")) + monkeypatch.setattr(response, "close", close) + monkeypatch.setattr(requests, "post", Mock(return_value=response)) - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.text_embedding.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError, match="broken status"), - ): + with pytest.raises(CredentialsValidateFailedError, match="401"): OAICompatEmbeddingModel([]).validate_credentials( "model", {"endpoint_url": "https://example.com/v1"} ) - response.close.assert_called_once_with() - - -def test_validate_credentials_wraps_credential_preparation_errors() -> None: - error_message = "lazy credentials failed" - sentinel = CredentialsValidateFailedError(error_message) - credentials = MagicMock() - credentials.get.side_effect = sentinel - - with pytest.raises(CredentialsValidateFailedError) as exc_info: - OAICompatEmbeddingModel([]).validate_credentials( - "model", - credentials, - ) - - assert exc_info.value is sentinel - - -def test_validate_credentials_accepts_duck_typed_response() -> None: - json_result = DuckEmbeddingResponse() - response = MagicMock(status_code=HTTPStatus.OK) - response.json.return_value = json_result - - with patch( - "dify_plugin.interfaces.model.openai_compatible.text_embedding.requests.post", - return_value=response, - ): - OAICompatEmbeddingModel([]).validate_credentials( - "model", - {"endpoint_url": "https://example.com/v1"}, - ) - - assert json_result.trace == ["contains"] - - -def test_validate_credentials_preserves_response_credential_errors() -> None: - sentinel = CredentialsValidateFailedError("response failed") - response = MagicMock(status_code=HTTPStatus.OK) - response.json.side_effect = sentinel - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.text_embedding.requests.post", - return_value=response, - ), - pytest.raises(CredentialsValidateFailedError) as exc_info, - ): - OAICompatEmbeddingModel([]).validate_credentials( - "model", - {"endpoint_url": "https://example.com/v1"}, - ) - - assert exc_info.value is sentinel - - -def test_validate_credentials_preserves_string_conversion_errors() -> None: - sentinel = RuntimeError() - response = MagicMock(status_code=HTTPStatus.OK) - response.json.side_effect = UnstringableResponseError(sentinel) - - with ( - patch( - "dify_plugin.interfaces.model.openai_compatible.text_embedding.requests.post", - return_value=response, - ), - pytest.raises(RuntimeError) as exc_info, - ): - OAICompatEmbeddingModel([]).validate_credentials( - "model", - {"endpoint_url": "https://example.com/v1"}, - ) - - assert exc_info.value is sentinel + close.assert_called_once_with() diff --git a/tests/interfaces/model/test_ai_model_parameter_rules.py b/tests/interfaces/model/test_ai_model_parameter_rules.py index c276887c..948b59d8 100644 --- a/tests/interfaces/model/test_ai_model_parameter_rules.py +++ b/tests/interfaces/model/test_ai_model_parameter_rules.py @@ -1,6 +1,6 @@ from collections.abc import Mapping from typing import ClassVar -from unittest.mock import MagicMock +from unittest.mock import Mock import pytest from pydantic import ConfigDict @@ -106,34 +106,61 @@ def test_parameter_rule_preserves_model_copy_update() -> None: assert schema.parameter_rules[0].max == 7 -def test_parameter_rule_template_keeps_legacy_value_error_boundary() -> None: - schema = CustomTemplateAIModel([]).get_customizable_model_schema_from_credentials( - "frozen", {} +def test_custom_model_schema_fills_zero_values_from_template( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = CustomTemplateAIModel([]) + schema = AIModelEntity( + model="custom", + model_type=ModelType.LLM, + model_properties={}, + parameter_rules=[ + ParameterRule.model_validate({ + "name": "top_p", + "use_template": "top_p", + "min": 0, + "max": 0, + "default": 0, + "precision": 0, + }) + ], + ) + monkeypatch.setattr( + model, "get_customizable_model_schema", Mock(return_value=schema) ) - assert schema is not None - assert isinstance(schema.parameter_rules[0], FrozenParameterRule) - assert schema.parameter_rules[0].precision == 0 + result = model.get_model_schema("custom", {"api_key": "key"}) + + assert result is not None + rule = result.parameter_rules[0] + assert (rule.min, rule.max, rule.default, rule.precision) == (0, 8, 1, 2) -def test_parameter_rule_template_copy_value_error_is_ignored( +def test_custom_model_schema_fills_missing_template_translations( monkeypatch: pytest.MonkeyPatch, ) -> None: - error_message = "lazy template failed" - template = MagicMock() - template.__bool__.return_value = True - template.copy.side_effect = ValueError(error_message) - monkeypatch.setitem( - PARAMETER_RULE_TEMPLATE, - DefaultParameterName.TOP_K, - template, + model = CustomTemplateAIModel([]) + schema = model.get_customizable_model_schema("copy", {}) + schema.parameter_rules[0].help = I18nObject(en_us="", zh_hans="") + monkeypatch.setattr( + model, "get_customizable_model_schema", Mock(return_value=schema) ) - rule = ParameterRule.model_validate({ - "name": "top_k", - "use_template": "top_k", - "label": I18nObject(en_us="Top K"), - "type": ParameterType.INT, - }) + result = model.get_model_schema("copy", {"api_key": "key"}) + + assert result is not None + help_text = result.parameter_rules[0].help + template_help = PARAMETER_RULE_TEMPLATE[DefaultParameterName.TOP_K]["help"] + assert help_text is not None + assert help_text.en_us == template_help["en_US"] + assert help_text.zh_hans == template_help["zh_Hans"] + + +def test_parameter_rule_template_keeps_legacy_value_error_boundary() -> None: + schema = CustomTemplateAIModel([]).get_customizable_model_schema_from_credentials( + "frozen", {} + ) - assert rule.name == "top_k" + assert schema is not None + assert isinstance(schema.parameter_rules[0], FrozenParameterRule) + assert schema.parameter_rules[0].precision == 0 diff --git a/tests/interfaces/model/test_ai_model_timing_context.py b/tests/interfaces/model/test_ai_model_timing_context.py index 0ed98838..7eecfa1a 100644 --- a/tests/interfaces/model/test_ai_model_timing_context.py +++ b/tests/interfaces/model/test_ai_model_timing_context.py @@ -1,112 +1,99 @@ -import time from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor +from threading import Barrier, Event import pytest +from dify_plugin.entities.model import ModelType from dify_plugin.errors.model import InvokeError from dify_plugin.interfaces.exec.ai_model import TimingContextRaceConditionError from dify_plugin.interfaces.model.ai_model import AIModel class MockAIModel(AIModel): - def validate_credentials(self, model: str, credentials: Mapping) -> None: - """Validate model credentials + model_type = ModelType.LLM - :param model: model name - :param credentials: model credentials - :return: None - """ + def validate_credentials(self, model: str, credentials: Mapping) -> None: + del model, credentials + @property def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]: - """Map model invoke error to unified error - - :return: Invoke error mapping - - Returns: - The return value. - """ return {} - def invoke(self) -> float: - """Invoke model""" - with self.timing_context(): - time.sleep(1) - return time.perf_counter() - self.started_at - - -class CleanupFailingAIModel(MockAIModel): - reject_cleanup = False - - def __setattr__(self, name: str, value: object) -> None: - if name == "started_at" and value == 0 and self.reject_cleanup: - msg = "cleanup failed" - raise ValueError(msg) - super().__setattr__(name, value) - - -def test_ai_model_timing_context_with_race_condition() -> None: - model = MockAIModel(model_schemas=[]) - concurrency = 2 - - def task(_: int) -> None: - """Task to be executed in thread pool""" - model.invoke() - - with ( - pytest.raises(TimingContextRaceConditionError), - ThreadPoolExecutor(concurrency) as pool, - ): - list(pool.map(task, range(concurrency))) - - -def test_ai_model_timing_context_multiple_sequential_uses() -> None: - """Check if multiple sequential uses of the timing context are correct""" - model = MockAIModel(model_schemas=[]) - - time_cost = model.invoke() - assert time_cost > 0 - assert time_cost < 1.5 - time_cost = model.invoke() - assert time_cost > 0 - assert time_cost < 1.5 +def test_ai_model_timing_context_rejects_overlapping_use() -> None: + model = MockAIModel([]) + entered = Event() + release = Event() + + def owner() -> float: + with model.timing_context(): + started_at = model.started_at + entered.set() + assert release.wait(timeout=5) + assert model.started_at == started_at + return started_at + + def contender() -> None: + with model.timing_context(): + pytest.fail("An active timing context must reject a second caller") + + with ThreadPoolExecutor(max_workers=2) as pool: + running = pool.submit(owner) + try: + assert entered.wait(timeout=5) + rejected = pool.submit(contender) + with pytest.raises(TimingContextRaceConditionError): + rejected.result(timeout=5) + finally: + release.set() + assert running.result(timeout=5) > 0 assert model.started_at == 0 -def test_ai_model_timing_context_resets_after_error() -> None: - model = MockAIModel(model_schemas=[]) - - def fail() -> None: - raise RuntimeError +def test_ai_model_timing_context_multiple_sequential_uses( + monkeypatch: pytest.MonkeyPatch, +) -> None: + timestamps = iter([10.0, 25.0]) + monkeypatch.setattr( + "dify_plugin.interfaces.model.ai_model.time.perf_counter", + lambda: next(timestamps), + ) + model = MockAIModel([]) - with pytest.raises(RuntimeError), model.timing_context(): - fail() + for expected in (10.0, 25.0): + with model.timing_context(): + assert model.started_at == expected + assert model.started_at == 0 - assert model.started_at == 0 +@pytest.mark.parametrize("error", [RuntimeError("failed"), GeneratorExit()]) +def test_ai_model_timing_context_resets_after_error(error: BaseException) -> None: + model = MockAIModel([]) -def test_ai_model_timing_cleanup_does_not_mask_body_error() -> None: - body_error = RuntimeError() - model = CleanupFailingAIModel(model_schemas=[]) - model.reject_cleanup = True + with pytest.raises(type(error)) as exc_info, model.timing_context(): + raise error - with pytest.raises(RuntimeError) as exc_info, model.timing_context(): - raise body_error + assert exc_info.value is error + assert model.started_at == 0 - assert exc_info.value is body_error +def test_ai_model_timing_context_is_independent_per_instance() -> None: + models = [MockAIModel([]) for _ in range(4)] + rendezvous = Barrier(len(models)) -def test_ai_model_timing_context_check_latency() -> None: - concurrency = 10 + def invoke(model: MockAIModel) -> float: + with model.timing_context(): + started_at = model.started_at + rendezvous.wait(timeout=5) + assert all(other.started_at > 0 for other in models) + rendezvous.wait(timeout=5) + return started_at - def task(_: int) -> None: - """Check if race condition is raised""" - model = MockAIModel(model_schemas=[]) - time_cost = model.invoke() - assert time_cost > 0 - assert time_cost < 1.5 + with ThreadPoolExecutor(max_workers=len(models)) as pool: + starts = list(pool.map(invoke, models)) - with ThreadPoolExecutor(concurrency) as pool: - list(pool.map(task, range(concurrency))) + assert len(starts) == len(models) + assert all(start > 0 for start in starts) + assert all(model.started_at == 0 for model in models) diff --git a/tests/interfaces/model/test_llm_parameters.py b/tests/interfaces/model/test_llm_parameters.py new file mode 100644 index 00000000..9559470d --- /dev/null +++ b/tests/interfaces/model/test_llm_parameters.py @@ -0,0 +1,170 @@ +import math +from unittest.mock import Mock + +import pytest + +from dify_plugin.entities.model import AIModelEntity, ModelType, ParameterRule +from dify_plugin.entities.model.llm import LLMResult, LLMUsage +from dify_plugin.entities.model.message import AssistantPromptMessage +from tests.interfaces.model.utils import StubLLM + + +@pytest.fixture +def llm() -> StubLLM: + rules = [ + { + "name": "max_tokens", + "type": "int", + "min": 1, + "max": 8, + "required": True, + "default": 2, + }, + {"name": "temperature", "type": "float", "min": 0, "max": 2, "precision": 2}, + {"name": "penalty", "type": "float", "min": -2, "max": 2}, + {"name": "count", "type": "float", "precision": 0}, + {"name": "enabled", "type": "boolean"}, + {"name": "mode", "type": "string", "options": ["fast", "accurate"]}, + {"name": "instructions", "type": "text"}, + {"name": "k", "use_template": "top_k", "min": 1, "max": 10}, + ] + return StubLLM([ + AIModelEntity( + model="model", + model_type=ModelType.LLM, + model_properties={}, + parameter_rules=[ParameterRule.model_validate(rule) for rule in rules], + ) + ]) + + +@pytest.mark.parametrize( + "parameters", + [ + {"max_tokens": 1}, + {"max_tokens": 8}, + {"max_tokens": True}, + {"temperature": 0}, + {"temperature": 2.0}, + {"temperature": 1.25}, + {"temperature": False}, + {"penalty": 0.125}, + {"count": 2.0}, + {"enabled": False}, + {"mode": "accurate"}, + {"instructions": ""}, + ], +) +def test_invoke_validates_and_forwards_model_parameters( + llm: StubLLM, + parameters: dict[str, object], + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = Mock( + return_value=LLMResult( + model="model", + message=AssistantPromptMessage(content="ok"), + usage=LLMUsage.empty_usage(), + ) + ) + monkeypatch.setattr(llm, "_invoke", provider) + supplied = {**parameters, "unrecognized": "drop"} + + chunks = list(llm.invoke("model", {}, [], model_parameters=supplied)) + + assert len(chunks) == 1 + assert chunks[0].delta.message.content == "ok" + assert provider.call_args.kwargs["model_parameters"] == { + "max_tokens": 2, + **parameters, + } + assert supplied == {**parameters, "unrecognized": "drop"} + assert llm.started_at == 0 + + +@pytest.mark.parametrize( + ("parameters", "expected"), + [ + ({}, {"max_tokens": 2}), + ({"max_tokens": None, "temperature": None}, {"max_tokens": 2}), + ({"top_k": 3}, {"max_tokens": 2, "k": 3}), + ({"top_k": 3, "k": 4}, {"max_tokens": 2, "k": 4}), + ({"top_k": 3, "k": None}, {"max_tokens": 2, "k": 3}), + ], +) +def test_model_parameter_defaults_and_aliases( + llm: StubLLM, + parameters: dict[str, object], + expected: dict[str, object], +) -> None: + assert ( + llm._validate_and_filter_model_parameters("model", parameters, {}) == expected + ) + + +@pytest.mark.parametrize( + ("name", "value", "message"), + [ + ("max_tokens", 0, "greater than or equal"), + ("max_tokens", 9, "less than or equal"), + ("max_tokens", 1.0, "should be int"), + ("max_tokens", "2", "should be int"), + ("temperature", -0.01, "greater than or equal"), + ("temperature", 2.01, "less than or equal"), + ("temperature", "1.2", "should be float"), + ("temperature", 1.234, "decimal places"), + ("count", 2.5, "should be int"), + ("enabled", 1, "should be bool"), + ("enabled", "false", "should be bool"), + ("mode", "unknown", "should be one of"), + ("mode", 1, "should be string"), + ("instructions", [], "should be string"), + ("top_k", 11, "less than or equal"), + ], +) +def test_invalid_model_parameters_fail_before_provider_invocation( + llm: StubLLM, + name: str, + value: object, + message: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + provider = Mock( + return_value=LLMResult( + model="model", + message=AssistantPromptMessage(content="ok"), + usage=LLMUsage.empty_usage(), + ) + ) + monkeypatch.setattr(llm, "_invoke", provider) + + with pytest.raises(ValueError, match=message): + list(llm.invoke("model", {}, [], model_parameters={name: value})) + + provider.assert_not_called() + assert llm.started_at == 0 + + +@pytest.mark.parametrize("default", [True, 0, 9]) +def test_required_model_parameter_keeps_provider_default( + llm: StubLLM, default: object +) -> None: + llm.model_schemas[0].parameter_rules[0].default = default + + assert llm._validate_and_filter_model_parameters("model", {}, {}) == { + "max_tokens": default + } + + +def test_required_model_parameter_without_default_fails(llm: StubLLM) -> None: + llm.model_schemas[0].parameter_rules[0].default = None + + with pytest.raises(ValueError, match="is required"): + llm._validate_and_filter_model_parameters("model", {}, {}) + + +def test_float_parameter_keeps_nan_acceptance(llm: StubLLM) -> None: + parameters = llm._validate_and_filter_model_parameters( + "model", {"penalty": float("nan")}, {} + ) + assert math.isnan(parameters["penalty"]) diff --git a/tests/interfaces/model/test_llm_stop_tokens.py b/tests/interfaces/model/test_llm_stop_tokens.py new file mode 100644 index 00000000..ee9d94f0 --- /dev/null +++ b/tests/interfaces/model/test_llm_stop_tokens.py @@ -0,0 +1,22 @@ +import pytest + +from tests.interfaces.model.utils import StubLLM + + +@pytest.mark.parametrize( + ("text", "stop", "expected"), + [ + ("hello.world", [r"\."], "hello"), + ("hello.world", ["."], ""), + ("answer[END]ignored", ["[END]"], "answer["), + ("answer\\ignored", [r"\\"], "answer"), + ("first\nsecond", [r"\n"], "first"), + ("firstsecond", ["", ""], "first"), + ("keep all of this", [""], "keep all of this"), + ("keep all of this", [], ""), + ], +) +def test_stop_sequences_preserve_regex_matching( + text: str, stop: list[str], expected: str +) -> None: + assert StubLLM([]).enforce_stop_tokens(text, stop) == expected diff --git a/tests/interfaces/model/test_llm_stream_errors.py b/tests/interfaces/model/test_llm_stream_errors.py index 4ab06c85..43711c84 100644 --- a/tests/interfaces/model/test_llm_stream_errors.py +++ b/tests/interfaces/model/test_llm_stream_errors.py @@ -2,7 +2,6 @@ from unittest.mock import patch import pytest -from requests.exceptions import ReadTimeout from dify_plugin.entities.model.llm import LLMResultChunk, LLMResultChunkDelta from dify_plugin.entities.model.message import AssistantPromptMessage @@ -11,16 +10,14 @@ InvokeError, InvokeRateLimitError, ) -from dify_plugin.interfaces.model.openai_compatible.llm import ( - OAICompatLargeLanguageModel, -) +from tests.interfaces.model.utils import StubLLM @pytest.mark.parametrize( ("error", "expected_type"), [ - (ReadTimeout("stream timed out"), InvokeConnectionError), - (ValueError("provider stream failed"), InvokeError), + (TimeoutError("stream timed out"), InvokeConnectionError), + (RuntimeError("provider stream failed"), InvokeError), (InvokeRateLimitError(description="rate limited"), InvokeRateLimitError), ], ) @@ -28,7 +25,7 @@ def test_invoke_maps_errors_after_stream_starts( error: Exception, expected_type: type[InvokeError], ) -> None: - llm = OAICompatLargeLanguageModel([]) + llm = StubLLM([]) first_chunk = LLMResultChunk( model="model", delta=LLMResultChunkDelta( @@ -36,6 +33,7 @@ def test_invoke_maps_errors_after_stream_starts( message=AssistantPromptMessage(content="\nA"), ), ) + expected_chunk = first_chunk.model_copy(deep=True) def response() -> Generator[LLMResultChunk, None, None]: yield first_chunk @@ -43,7 +41,7 @@ def response() -> Generator[LLMResultChunk, None, None]: with patch.object(llm, "_invoke", return_value=response()): stream = llm.invoke(model="model", credentials={}, prompt_messages=[]) - assert next(stream) is first_chunk + assert next(stream) == expected_chunk assert llm.started_at > 0 with pytest.raises(expected_type) as exc_info: @@ -57,7 +55,7 @@ def response() -> Generator[LLMResultChunk, None, None]: def test_invoke_closes_provider_stream() -> None: - llm = OAICompatLargeLanguageModel([]) + llm = StubLLM([]) closed = False def response() -> Generator[LLMResultChunk, None, None]: @@ -80,3 +78,53 @@ def response() -> Generator[LLMResultChunk, None, None]: assert closed assert llm.started_at == 0 + + +@pytest.mark.parametrize( + ("error", "expected_type"), + [ + (TimeoutError("request timed out"), InvokeConnectionError), + (RuntimeError("provider failed"), InvokeError), + (InvokeRateLimitError(description="rate limited"), InvokeRateLimitError), + ], +) +def test_invoke_maps_errors_before_first_chunk( + error: Exception, + expected_type: type[InvokeError], +) -> None: + llm = StubLLM([]) + + with ( + patch.object(llm, "_invoke", side_effect=error), + pytest.raises(expected_type) as exc_info, + ): + list(llm.invoke(model="model", credentials={}, prompt_messages=[])) + + if isinstance(error, InvokeError): + assert exc_info.value is error + else: + assert exc_info.value.__cause__ is error + assert llm.started_at == 0 + + +def test_empty_llm_stream_finishes_and_resets_timing() -> None: + llm = StubLLM([]) + + def response() -> Generator[LLMResultChunk]: + yield from () + + with patch.object(llm, "_invoke", return_value=response()): + assert list(llm.invoke(model="model", credentials={}, prompt_messages=[])) == [] + + assert llm.started_at == 0 + + +def test_closing_unstarted_llm_stream_does_not_invoke_provider() -> None: + llm = StubLLM([]) + + with patch.object(llm, "_invoke") as provider: + stream = llm.invoke(model="model", credentials={}, prompt_messages=[]) + stream.close() + + provider.assert_not_called() + assert llm.started_at == 0 diff --git a/tests/interfaces/model/test_llm_structured_output.py b/tests/interfaces/model/test_llm_structured_output.py new file mode 100644 index 00000000..64bcdeea --- /dev/null +++ b/tests/interfaces/model/test_llm_structured_output.py @@ -0,0 +1,683 @@ +from collections.abc import Generator +from copy import deepcopy +from types import MappingProxyType +from unittest.mock import Mock + +import pytest + +from dify_plugin.entities.model import AIModelEntity, ModelType, ParameterRule +from dify_plugin.entities.model.llm import ( + LLMResult, + LLMResultChunk, + LLMResultChunkDelta, + LLMUsage, +) +from dify_plugin.entities.model.message import ( + AssistantPromptMessage, + ImagePromptMessageContent, + PromptMessage, + SystemPromptMessage, + TextPromptMessageContent, + UserPromptMessage, +) +from dify_plugin.errors.model import InvokeConnectionError, InvokeError +from tests.interfaces.model.utils import StubLLM + + +@pytest.fixture +def llm() -> StubLLM: + return StubLLM([ + AIModelEntity( + model="model", + model_type=ModelType.LLM, + model_properties={}, + parameter_rules=[ + ParameterRule.model_validate({ + "name": "response_format", + "type": "string", + "options": ["JSON", "XML"], + }) + ], + ) + ]) + + +def _chunk(content: str) -> LLMResultChunk: + return LLMResultChunk( + model="model", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content=content), + ), + ) + + +@pytest.mark.parametrize("response_format", ["JSON", "XML"]) +@pytest.mark.parametrize("split_at", [1, 2, 3, 6, 12]) +def test_structured_stream_strips_fragmented_fences_and_preserves_usage( + llm: StubLLM, + response_format: str, + split_at: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + body = '{"answer":"ok"}' if response_format == "JSON" else "ok" + text = f"```{response_format.lower()}\n{body}\n```ignored" + usage = LLMUsage.empty_usage() + usage.prompt_tokens = 7 + usage.completion_tokens = 3 + usage.total_tokens = 10 + terminal = _chunk("") + terminal.delta.usage = usage + terminal.delta.finish_reason = "stop" + terminal.system_fingerprint = "provider-fingerprint" + expected_terminal = terminal.model_copy(deep=True) + + def response() -> Generator[LLMResultChunk]: + yield _chunk(text[:split_at]) + yield _chunk(text[split_at:]) + yield terminal + + monkeypatch.setattr(llm, "_invoke", Mock(return_value=response())) + chunks = list( + llm.invoke( + "model", + {}, + [], + model_parameters={"response_format": response_format}, + ) + ) + + assert ( + "".join(str(chunk.delta.message.content or "") for chunk in chunks) + == body + "\n" + ) + assert chunks[-1] == expected_terminal + assert chunks[-1].delta.finish_reason == "stop" + assert chunks[-1].system_fingerprint == "provider-fingerprint" + assert llm.started_at == 0 + + +@pytest.mark.parametrize("leading_empty", [False, True]) +def test_structured_stream_accepts_plain_json_and_empty_role_chunks( + llm: StubLLM, + leading_empty: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def response() -> Generator[LLMResultChunk]: + if leading_empty: + yield _chunk("") + yield _chunk('{"answer":') + yield _chunk('"ok"}') + + monkeypatch.setattr(llm, "_invoke", Mock(return_value=response())) + chunks = list( + llm.invoke( + "model", + {}, + [], + model_parameters={"response_format": "JSON"}, + ) + ) + + assert ( + "".join(str(chunk.delta.message.content or "") for chunk in chunks) + == '{"answer":"ok"}' + ) + assert llm.started_at == 0 + + +@pytest.mark.parametrize("text_parts", [False, True]) +def test_structured_output_prepares_system_and_user_instructions( + llm: StubLLM, + text_parts: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + user_content = ( + [TextPromptMessageContent(data="question")] if text_parts else "question" + ) + prompts = [ + SystemPromptMessage(content="Return the answer"), + UserPromptMessage(content=user_content), + ] + provider = Mock(return_value=(_chunk("{}") for _ in range(1))) + monkeypatch.setattr(llm, "_invoke", provider) + + chunks = list( + llm.invoke( + "model", + {}, + prompts, + model_parameters={"response_format": "JSON"}, + stop=["END"], + user="user-1", + ) + ) + + assert "".join(str(chunk.delta.message.content or "") for chunk in chunks) == "{}" + forwarded = provider.call_args.kwargs + assert forwarded["model_parameters"] == {} + assert forwarded["stop"] == ["END", "\n```", "```\n"] + assert forwarded["user"] == "user-1" + assert "Return the answer" in forwarded["prompt_messages"][0].content + content = forwarded["prompt_messages"][-1].content + text = content[-1].data if text_parts else content + assert text == "question\n```JSON\n" + + +def test_empty_structured_stream_keeps_provider_error( + llm: StubLLM, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(llm, "_invoke", Mock(return_value=(_chunk("") for _ in ()))) + + with pytest.raises(InvokeError) as exc_info: + list( + llm.invoke( + "model", + {}, + [], + model_parameters={"response_format": "JSON"}, + ) + ) + assert isinstance(exc_info.value.__cause__, StopIteration) + assert llm.started_at == 0 + + +@pytest.mark.parametrize("first_content", ["```json\n{}", "{}"]) +@pytest.mark.parametrize("consume_first_chunk", [False, True]) +def test_closing_structured_stream_closes_provider( + llm: StubLLM, + first_content: str, + consume_first_chunk: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + closed = [] + + def response() -> Generator[LLMResultChunk]: + try: + yield _chunk(first_content) + pytest.fail("Closing the public stream must not consume more provider data") + finally: + closed.append(True) + + provider_stream = response() + monkeypatch.setattr(llm, "_invoke", Mock(return_value=provider_stream)) + invoke = llm.invoke if consume_first_chunk else llm._code_block_mode_wrapper + stream = invoke("model", {}, [], model_parameters={"response_format": "JSON"}) + assert isinstance(stream, Generator) + if consume_first_chunk: + next(stream) + stream.close() + + assert closed == [True] + assert llm.started_at == 0 + + +def test_structured_stream_maps_provider_failure_and_closes( + llm: StubLLM, + monkeypatch: pytest.MonkeyPatch, +) -> None: + error = TimeoutError("provider stalled") + closed = [] + + def response() -> Generator[LLMResultChunk]: + try: + yield _chunk("```json\n{}") + raise error + finally: + closed.append(True) + + monkeypatch.setattr(llm, "_invoke", Mock(return_value=response())) + + with pytest.raises(InvokeConnectionError) as exc_info: + list(llm.invoke("model", {}, [], model_parameters={"response_format": "JSON"})) + + assert exc_info.value.__cause__ is error + assert closed == [True] + assert llm.started_at == 0 + + +@pytest.mark.parametrize("fenced", [False, True]) +@pytest.mark.parametrize("leading_empty", [False, True]) +def test_structured_stream_preserves_backticks_inside_json_strings( + llm: StubLLM, fenced: bool, leading_empty: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + body = '{"text":"one`two``three"}' + source = f"```json\n{body}\n```" if fenced else body + + def response() -> Generator[LLMResultChunk]: + if leading_empty: + yield _chunk("") + for character in source: + yield _chunk(character) + + monkeypatch.setattr(llm, "_invoke", Mock(return_value=response())) + + chunks = list( + llm.invoke("model", {}, [], model_parameters={"response_format": "JSON"}) + ) + + expected = body + "\n" if fenced else body + assert ( + "".join(str(chunk.delta.message.content or "") for chunk in chunks) == expected + ) + assert llm.started_at == 0 + + +def test_structured_output_preserves_image_content_and_complete_provider_result( + llm: StubLLM, monkeypatch: pytest.MonkeyPatch +) -> None: + media = ImagePromptMessageContent( + format="png", mime_type="image/png", url="https://example.com/image.png" + ) + response = LLMResult( + model="model", + message=AssistantPromptMessage(content='{"answer":42}'), + usage=LLMUsage.empty_usage(), + ) + response.usage.prompt_tokens = 11 + response.usage.completion_tokens = 2 + response.usage.total_tokens = 13 + expected_response = response.model_copy(deep=True) + provider = Mock(return_value=response) + monkeypatch.setattr(llm, "_invoke", provider) + + chunks = list( + llm.invoke( + "model", + {}, + [UserPromptMessage(content=[media])], + model_parameters={"response_format": "JSON"}, + stream=False, + ) + ) + + assert len(chunks) == 1 + assert chunks[0].delta.message == expected_response.message + assert chunks[0].delta.usage == expected_response.usage + forwarded = provider.call_args.kwargs + assert forwarded["stream"] is False + assert forwarded["prompt_messages"][-1].content == [media] + assert "valid JSON object" in forwarded["prompt_messages"][0].content + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ('```json\n{"answer":42}\n```\nexplanation', '{"answer":42}\nexplanation'), + ('{"text":"use ``` to mark code"}', '{"text":"use to mark code"}'), + ( + '```json\n{"text":"use ``` to mark code"}\n```\nexplanation', + '{"text":"use to mark code"}\nexplanation', + ), + ], + ids=["empty-before-fence", "inline-backticks", "fenced-inline-backticks"], +) +@pytest.mark.parametrize("chunk_size", [1, 2, 7, 100]) +def test_structured_stream_keeps_fence_handling_after_empty_first_chunk( + llm: StubLLM, + source: str, + expected: str, + chunk_size: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + leading = _chunk("") + terminal = _chunk("") + terminal.delta.usage = LLMUsage.empty_usage() + terminal.delta.finish_reason = "stop" + terminal.system_fingerprint = "provider-revision" + expected_leading = leading.model_copy(deep=True) + expected_terminal = terminal.model_copy(deep=True) + + def response() -> Generator[LLMResultChunk]: + yield leading + for offset in range(0, len(source), chunk_size): + text = source[offset : offset + chunk_size] + yield _chunk(text) + yield terminal + + monkeypatch.setattr(llm, "_invoke", Mock(return_value=response())) + + chunks = list( + llm.invoke("model", {}, [], model_parameters={"response_format": "JSON"}) + ) + + assert ( + "".join(chunk.delta.message.get_text_content() for chunk in chunks) == expected + ) + assert chunks[0] == expected_leading + assert chunks[-1] == expected_terminal + assert chunks[-1].delta.finish_reason == "stop" + assert chunks[-1].system_fingerprint == "provider-revision" + assert llm.started_at == 0 + + +@pytest.mark.parametrize("text_parts", [False, True]) +def test_structured_invocations_keep_input_updates_when_reused( + llm: StubLLM, text_parts: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + prompts = [ + SystemPromptMessage(content="Return the answer"), + UserPromptMessage( + content=[TextPromptMessageContent(data="question")] + if text_parts + else "question" + ), + ] + stop = ["END"] + calls = [] + + def invoke(**kwargs: object) -> LLMResult: + calls.append(deepcopy(kwargs)) + return LLMResult( + model="model", + message=AssistantPromptMessage(content="{}"), + usage=LLMUsage.empty_usage(), + ) + + monkeypatch.setattr(llm, "_invoke", invoke) + for _ in range(2): + list(llm.invoke("model", {}, prompts, {"response_format": "JSON"}, stop=stop)) + + instruction_prefix = ( + "You should always follow the instructions and output a valid JSON object.\n" + "The structure of the JSON object you can found in the instructions, " + 'use {"answer": "$your_answer"} as the default structure\n' + "if you are not sure about the structure.\n\n\n" + ) + first_system = instruction_prefix + "Return the answer\n\n" + assert calls[0]["prompt_messages"][0].content == first_system + assert calls[1]["prompt_messages"][0].content == ( + instruction_prefix + first_system + "\n\n" + ) + assert calls[0]["prompt_messages"][-1].get_text_content() == "question\n```JSON\n" + assert ( + calls[1]["prompt_messages"][-1].get_text_content() + == "question\n```JSON\n\n```JSON\n" + ) + assert calls[0]["stop"] == ["END", "\n```", "```\n"] + assert calls[1]["stop"] == ["END", "\n```", "```\n", "\n```", "```\n"] + assert prompts == calls[1]["prompt_messages"] + assert stop == calls[1]["stop"] + + +def test_structured_wrapper_keeps_response_format_pop_when_parameters_are_reused( + llm: StubLLM, monkeypatch: pytest.MonkeyPatch +) -> None: + prompts = [UserPromptMessage(content="question")] + parameters = {"response_format": "JSON"} + stop = ["END"] + calls = [] + + def invoke(**kwargs: object) -> LLMResult: + calls.append(deepcopy(kwargs)) + return LLMResult( + model="model", + message=AssistantPromptMessage(content="{}"), + usage=LLMUsage.empty_usage(), + ) + + monkeypatch.setattr(llm, "_invoke", invoke) + for _ in range(2): + llm._code_block_mode_wrapper( + "model", {}, prompts, parameters, stop=stop, stream=False + ) + + assert parameters == {} + assert calls[0] == calls[1] + assert calls[0]["model_parameters"] == {} + assert prompts[-1].content == "question\n```JSON\n" + assert stop == ["END", "\n```", "```\n"] + + +@pytest.mark.parametrize("text_parts", [False, True]) +def test_structured_invocations_preserve_repeated_prompt_references( + llm: StubLLM, text_parts: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + text = TextPromptMessageContent(data="question") + prompt = UserPromptMessage(content=[text] if text_parts else "question") + prompts = [prompt, prompt] + provider = Mock( + return_value=LLMResult( + model="model", + message=AssistantPromptMessage(content="{}"), + usage=LLMUsage.empty_usage(), + ) + ) + monkeypatch.setattr(llm, "_invoke", provider) + + for expected in ["question\n```JSON\n", "question\n```JSON\n\n```JSON\n"]: + list( + llm.invoke("model", {}, prompts, {"response_format": "JSON"}, stream=False) + ) + forwarded = provider.call_args.kwargs["prompt_messages"] + assert [message.get_text_content() for message in forwarded[1:]] == [ + expected, + expected, + ] + assert forwarded[1] is forwarded[2] is prompt + assert prompts is forwarded + assert prompt.get_text_content() == expected + assert text.data == (expected if text_parts else "question") + + +@pytest.mark.parametrize("shared_content_list", [False, True]) +def test_structured_invocations_preserve_shared_text_references( + llm: StubLLM, shared_content_list: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + shared = TextPromptMessageContent(data="question") + independent = TextPromptMessageContent(data="question") + previous = AssistantPromptMessage(content=[shared]) + prompt = UserPromptMessage(content=[independent, shared, shared]) + if shared_content_list: + previous.content = prompt.content + prompts = [previous, prompt] + provider = Mock( + return_value=LLMResult( + model="model", + message=AssistantPromptMessage(content="{}"), + usage=LLMUsage.empty_usage(), + ) + ) + monkeypatch.setattr(llm, "_invoke", provider) + + for expected in ["question\n```JSON\n", "question\n```JSON\n\n```JSON\n"]: + list( + llm.invoke("model", {}, prompts, {"response_format": "JSON"}, stream=False) + ) + forwarded = provider.call_args.kwargs["prompt_messages"] + first, last = forwarded[1:] + assert [part.data for part in last.content] == ["question", expected, expected] + assert first.content[-1] is last.content[-1] + assert last.content[1] is last.content[2] + assert (first.content is last.content) is shared_content_list + assert last.content[0] is independent + assert shared.data == expected + assert independent.data == "question" + assert previous.content[-1] is prompt.content[-1] is shared + + +@pytest.mark.parametrize("text_parts", [False, True]) +def test_structured_invocation_preserves_provider_prompt_subclasses( + llm: StubLLM, text_parts: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + class ProviderPrompt(UserPromptMessage): + provider_state: object + + class ProviderText(TextPromptMessageContent): + provider_state: object + + state = MappingProxyType({"tenant": "test"}) + text = ProviderText(data="question", provider_state=state) + prompt = ProviderPrompt( + content=[text] if text_parts else "question", provider_state=state + ) + provider = Mock( + return_value=LLMResult( + model="model", + message=AssistantPromptMessage(content="{}"), + usage=LLMUsage.empty_usage(), + ) + ) + monkeypatch.setattr(llm, "_invoke", provider) + + chunks = list( + llm.invoke("model", {}, [prompt], {"response_format": "JSON"}, stream=False) + ) + + assert chunks[0].delta.message.content == "{}" + forwarded = provider.call_args.kwargs["prompt_messages"][-1] + assert forwarded is prompt + assert isinstance(forwarded, ProviderPrompt) + assert forwarded.provider_state is state + if text_parts: + assert isinstance(forwarded.content[0], ProviderText) + assert forwarded.content[0].provider_state is state + assert forwarded.content[0].data == "question\n```JSON\n" + assert prompt.content == [text] + assert text.data == "question\n```JSON\n" + else: + assert forwarded.content == "question\n```JSON\n" + assert prompt.content == "question\n```JSON\n" + + +@pytest.mark.parametrize( + ("source", "expected", "response_format", "finish_reason"), + [ + ("\n \t```json\n{}\n```", "\n \t{}\n", "JSON", "stop"), + (" None: + terminal = _chunk("") + terminal.delta.finish_reason = finish_reason + terminal.delta.usage = LLMUsage.empty_usage() + + def response() -> Generator[LLMResultChunk]: + for character in source: + yield _chunk(character) + yield terminal + + monkeypatch.setattr(llm, "_invoke", Mock(return_value=response())) + chunks = list(llm.invoke("model", {}, [], {"response_format": response_format})) + + assert ( + "".join(chunk.delta.message.get_text_content() for chunk in chunks) == expected + ) + assert [chunk.delta.finish_reason for chunk in chunks if chunk.delta.usage] == [ + finish_reason + ] + assert llm.started_at == 0 + + +@pytest.mark.parametrize("fenced", [False, True]) +@pytest.mark.parametrize("chunk_size", [1, 7]) +@pytest.mark.parametrize( + ("body", "plain_output", "fenced_output"), + [ + ( + "", + "", + "ok', + 'ok', + '", + "", + "