mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 22:13:01 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
638 lines
24 KiB
Python
638 lines
24 KiB
Python
"""Unit tests for native/extracted document passing to workflow AGENT nodes.
|
|
|
|
These exercise the engine's ``_materialize_node_attachments`` policy and the
|
|
``_execute_agent_node`` wiring with a fake artifacts repo (no live DB / storage /
|
|
sandbox / LLM). The run-scoped authz gate, the auto/native/extract policy matrix,
|
|
the caps, and the no-selection regression path are all covered here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import uuid
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict
|
|
|
|
import pytest
|
|
|
|
from docsgpt.agents.workflows.node_agent import WorkflowNodeAgentFactory
|
|
from docsgpt.agents.workflows.schemas import (
|
|
NodeType,
|
|
Workflow,
|
|
WorkflowGraph,
|
|
WorkflowNode,
|
|
)
|
|
from docsgpt.agents.workflows.workflow_engine import WorkflowEngine
|
|
from docsgpt.core.settings import settings
|
|
from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository
|
|
|
|
RUN_ID = "11111111-1111-1111-1111-111111111111"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures / fakes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _patch_repo(monkeypatch, artifacts: Dict[str, Dict[str, Any]], run_id: str = RUN_ID) -> None:
|
|
"""Patch db_readonly + the three repo methods the engine calls so it reads the fake data.
|
|
|
|
The repo CLASS is left intact (transitive ``from ... import ArtifactsRepository``
|
|
bindings are untouched); only the methods the resolver uses are redirected, which
|
|
monkeypatch reverts cleanly on teardown.
|
|
"""
|
|
# 1-based position order = insertion order of ids belonging to this run.
|
|
positions = [aid for aid, a in artifacts.items() if a.get("run_id") == run_id]
|
|
|
|
def _at_position(self, n, *, conversation_id=None, workflow_run_id=None):
|
|
if workflow_run_id != run_id or not isinstance(n, int) or n < 1 or n > len(positions):
|
|
return None
|
|
return positions[n - 1]
|
|
|
|
def _in_parent(self, artifact_id, *, conversation_id=None, workflow_run_id=None):
|
|
art = artifacts.get(artifact_id)
|
|
if art is None or art.get("run_id") != workflow_run_id:
|
|
return None
|
|
return {"current_version": art["current_version"], "title": art.get("title")}
|
|
|
|
def _version(self, artifact_id, version):
|
|
art = artifacts.get(artifact_id)
|
|
return art["versions"].get(version) if art is not None else None
|
|
|
|
@contextlib.contextmanager
|
|
def _fake_db_readonly():
|
|
yield object()
|
|
|
|
monkeypatch.setattr("docsgpt.storage.db.session.db_readonly", _fake_db_readonly)
|
|
monkeypatch.setattr(ArtifactsRepository, "__init__", lambda self, conn=None: None)
|
|
monkeypatch.setattr(ArtifactsRepository, "artifact_id_at_position", _at_position)
|
|
monkeypatch.setattr(ArtifactsRepository, "get_artifact_in_parent", _in_parent)
|
|
monkeypatch.setattr(ArtifactsRepository, "get_version", _version)
|
|
|
|
|
|
def _artifact(run_id, mime, *, filename="doc", size=10, version=1, storage_path=None):
|
|
"""Build a fake artifact record with a single current version."""
|
|
aid = str(uuid.uuid4())
|
|
return aid, {
|
|
"run_id": run_id,
|
|
"current_version": version,
|
|
"title": filename,
|
|
"versions": {
|
|
version: {
|
|
"storage_path": storage_path or f"store/{aid}",
|
|
"mime_type": mime,
|
|
"filename": filename,
|
|
"size": size,
|
|
}
|
|
},
|
|
}
|
|
|
|
|
|
def _engine(monkeypatch, run_id: str = RUN_ID) -> WorkflowEngine:
|
|
graph = WorkflowGraph(workflow=Workflow(name="Attach Test"), nodes=[], edges=[])
|
|
agent = SimpleNamespace(
|
|
endpoint="stream",
|
|
llm_name="openai",
|
|
model_id="gpt-4o-mini",
|
|
api_key="test-key",
|
|
chat_history=[],
|
|
user="user-1",
|
|
decoded_token={"sub": "user-1"},
|
|
)
|
|
return WorkflowEngine(graph, agent, workflow_run_id=run_id)
|
|
|
|
|
|
# Provider supported-types lists, mirroring what ``llm.get_supported_attachment_types()``
|
|
# returns -- the engine now decides native-eligibility from this same source.
|
|
VISION_TYPES = ["image/png", "image/jpeg"]
|
|
TEXT_TYPES: list = []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_input_artifact_ids: tokens, lists, refs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_star_token_expands_to_all_input_documents(monkeypatch):
|
|
"""``*`` expands to every ref in state['input_documents']."""
|
|
eng = _engine(monkeypatch)
|
|
eng.state["input_documents"] = [
|
|
{"artifact_id": "aaa"}, {"artifact_id": "bbb"}, {"not": "a ref"},
|
|
]
|
|
assert eng._resolve_input_artifact_ids(["*"]) == ["aaa", "bbb"]
|
|
assert eng._resolve_input_artifact_ids(["input_documents"]) == ["aaa", "bbb"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_state_var_single_ref_and_list_of_refs(monkeypatch):
|
|
"""A state var may hold a single ref dict or a list of ref dicts."""
|
|
eng = _engine(monkeypatch)
|
|
eng.state["one"] = {"artifact_id": "x1"}
|
|
eng.state["many"] = [{"artifact_id": "y1"}, {"artifact_id": "y2"}]
|
|
assert eng._resolve_input_artifact_ids(["one"]) == ["x1"]
|
|
assert eng._resolve_input_artifact_ids(["many"]) == ["y1", "y2"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_raw_id_and_short_ref_pass_through(monkeypatch):
|
|
"""A raw id or short ref that is not a state key passes through verbatim."""
|
|
eng = _engine(monkeypatch)
|
|
assert eng._resolve_input_artifact_ids(["A1", "abc-id"]) == ["A1", "abc-id"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _materialize_node_attachments: policy matrix
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _node_config(**kwargs):
|
|
from docsgpt.agents.workflows.schemas import AgentNodeConfig
|
|
|
|
return AgentNodeConfig(**kwargs)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_auto_image_goes_native_for_vision_model(monkeypatch):
|
|
"""auto + image mime + vision model -> a native attachment {id, mime_type, path}."""
|
|
aid, rec = _artifact(RUN_ID, "image/png", filename="chart.png")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
assert out == [{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_auto_pdf_native_when_model_supports_images(monkeypatch):
|
|
"""auto + PDF + vision (image) model -> native (synthetic PDF-to-image path)."""
|
|
aid, rec = _artifact(RUN_ID, "application/pdf", filename="r.pdf")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
assert out[0]["mime_type"] == "application/pdf"
|
|
assert "path" in out[0] and "content" not in out[0]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_auto_text_only_model_extracts_to_content(monkeypatch):
|
|
"""auto + text-only model -> extracted/inlined text content, non-native mime."""
|
|
aid, rec = _artifact(RUN_ID, "text/plain", filename="notes.txt")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
monkeypatch.setattr(
|
|
"docsgpt.storage.storage_creator.StorageCreator.get_storage",
|
|
staticmethod(lambda: _FakeStorage(b"hello world")),
|
|
)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES)
|
|
|
|
assert out == [{"id": aid, "mime_type": "text/plain", "content": "hello world"}]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_extract_always_inlines_text_even_for_vision_model(monkeypatch):
|
|
"""file_passing=extract inlines text even when the model could take it natively."""
|
|
aid, rec = _artifact(RUN_ID, "text/markdown", filename="r.md")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
monkeypatch.setattr(
|
|
"docsgpt.storage.storage_creator.StorageCreator.get_storage",
|
|
staticmethod(lambda: _FakeStorage(b"# Title")),
|
|
)
|
|
cfg = _node_config(input_documents=[aid], file_passing="extract")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
assert out == [{"id": aid, "mime_type": "text/plain", "content": "# Title"}]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_native_on_unsupported_mime_raises(monkeypatch):
|
|
"""file_passing=native on a mime the model can't read raises a clear error."""
|
|
aid, rec = _artifact(RUN_ID, "application/pdf", filename="r.pdf")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
cfg = _node_config(input_documents=[aid], file_passing="native", model_id="gpt-4o-mini")
|
|
|
|
with pytest.raises(ValueError, match="cannot read"):
|
|
eng._materialize_node_attachments(cfg, "MyNode", TEXT_TYPES)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_extract_non_text_uses_parse_worker(monkeypatch):
|
|
"""A non-text mime under extract routes through the parsing-worker path (no sandbox)."""
|
|
aid, rec = _artifact(RUN_ID, "application/vnd.openxmlformats", filename="r.docx")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
monkeypatch.setattr(
|
|
WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None, deadline=None: "EXTRACTED MD"
|
|
)
|
|
cfg = _node_config(input_documents=[aid], file_passing="extract")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES)
|
|
|
|
assert out == [{"id": aid, "mime_type": "text/plain", "content": "EXTRACTED MD"}]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Security: cross-run / forged refs are rejected, never attached
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_cross_run_artifact_is_rejected(monkeypatch):
|
|
"""An artifact belonging to a different run yields None at the gate -> ValueError."""
|
|
aid, rec = _artifact("99999999-9999-9999-9999-999999999999", "image/png")
|
|
_patch_repo(monkeypatch, {aid: rec}) # run scope is RUN_ID; artifact is in another run
|
|
eng = _engine(monkeypatch)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
with pytest.raises(ValueError, match="not found in this run"):
|
|
eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_forged_uuid_is_rejected(monkeypatch):
|
|
"""A forged uuid never matches the run-scoped gate -> ValueError, nothing attached."""
|
|
aid, rec = _artifact(RUN_ID, "image/png")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
forged = str(uuid.uuid4())
|
|
cfg = _node_config(input_documents=[forged], file_passing="auto")
|
|
|
|
with pytest.raises(ValueError, match="not found in this run"):
|
|
eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Caps
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_native_file_cap_bounds_native_then_extracts(monkeypatch):
|
|
"""More than the native cap: the first N go native, the rest are extracted."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.settings.settings.WORKFLOW_NODE_NATIVE_MAX_FILES", 2, raising=False
|
|
)
|
|
artifacts = {}
|
|
ids = []
|
|
for i in range(4):
|
|
aid, rec = _artifact(RUN_ID, "image/png", filename=f"i{i}.png")
|
|
artifacts[aid] = rec
|
|
ids.append(aid)
|
|
_patch_repo(monkeypatch, artifacts)
|
|
eng = _engine(monkeypatch)
|
|
monkeypatch.setattr(
|
|
"docsgpt.storage.storage_creator.StorageCreator.get_storage",
|
|
staticmethod(lambda: _FakeStorage(b"img-bytes")),
|
|
)
|
|
# extract of a non-text image routes through the parsing worker; stub it so it returns text.
|
|
monkeypatch.setattr(
|
|
WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None, deadline=None: "fallback text"
|
|
)
|
|
cfg = _node_config(input_documents=ids, file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
native = [a for a in out if "path" in a]
|
|
extracted = [a for a in out if "content" in a]
|
|
assert len(native) == 2
|
|
assert len(extracted) == 2
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_oversize_file_is_skipped(monkeypatch):
|
|
"""A file past the per-file byte ceiling is dropped, not attached."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.settings.settings.SANDBOX_MAX_INPUT_BYTES", 5, raising=False
|
|
)
|
|
aid, rec = _artifact(RUN_ID, "image/png", size=999)
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
assert out == []
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_oversize_text_skipped_by_post_read_guard_when_size_missing(monkeypatch):
|
|
"""A NULL/missing version size skips the pre-read cap; the post-read byte guard still drops it."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.settings.settings.SANDBOX_MAX_INPUT_BYTES", 5, raising=False
|
|
)
|
|
# size=None bypasses the ``isinstance(size, int)`` pre-read check; the actual
|
|
# bytes (longer than the 5-byte cap) must be rejected after reading.
|
|
aid, rec = _artifact(RUN_ID, "text/plain", filename="big.txt", size=None)
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
monkeypatch.setattr(
|
|
"docsgpt.storage.storage_creator.StorageCreator.get_storage",
|
|
staticmethod(lambda: _FakeStorage(b"way over the cap")),
|
|
)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES)
|
|
|
|
assert out == []
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_large_under_cap_text_is_windowed_not_inlined_whole(monkeypatch):
|
|
"""A large-but-under-cap text file is bounded to a head+tail window, not inlined whole."""
|
|
from docsgpt.parser.document_reader import _TEXT_MAX_BYTES as _MARKDOWN_MAX_BYTES
|
|
|
|
big_text = ("A" * (_MARKDOWN_MAX_BYTES * 3)).encode("utf-8")
|
|
aid, rec = _artifact(RUN_ID, "text/plain", filename="notes.txt", size=len(big_text))
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
monkeypatch.setattr(
|
|
"docsgpt.storage.storage_creator.StorageCreator.get_storage",
|
|
staticmethod(lambda: _FakeStorage(big_text)),
|
|
)
|
|
cfg = _node_config(input_documents=[aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", TEXT_TYPES)
|
|
|
|
assert len(out) == 1
|
|
content = out[0]["content"]
|
|
assert "...[truncated" in content
|
|
# Bounded to roughly the window, far below the whole 3x payload.
|
|
assert len(content.encode("utf-8")) < _MARKDOWN_MAX_BYTES + 200
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_duplicate_refs_attach_once(monkeypatch):
|
|
"""A duplicated/expanded ref resolves once so the same artifact is not attached twice."""
|
|
aid, rec = _artifact(RUN_ID, "image/png", filename="dup.png")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
eng.state["input_documents"] = [{"artifact_id": aid}]
|
|
cfg = _node_config(input_documents=["*", aid], file_passing="auto")
|
|
|
|
out = eng._materialize_node_attachments(cfg, "Node", VISION_TYPES)
|
|
|
|
assert out == [{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Size-scaled parse window
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _capture_parse(monkeypatch) -> dict:
|
|
"""Patch parse_document.apply_async so nothing touches a broker; capture the call."""
|
|
import docsgpt.api.user.tasks as tasks
|
|
|
|
captured: Dict[str, Any] = {}
|
|
|
|
class _R:
|
|
def get(self, timeout=None, disable_sync_subtasks=True):
|
|
captured["timeout"] = timeout
|
|
return {"status": "ok", "content": "PARSED"}
|
|
|
|
def _apply_async(args=None, queue=None, **kw):
|
|
captured["kwargs"] = kw
|
|
return _R()
|
|
|
|
monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
|
|
return captured
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_parse_window_scales_with_the_document_size(monkeypatch):
|
|
"""The version's ``size`` reaches the parse: a longer await AND matching Celery limits."""
|
|
from docsgpt.api.user.tasks import parse_task_time_limits, parse_timeout_for_size
|
|
|
|
size = 8 * 1024 * 1024
|
|
aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf", size=size)
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
captured = _capture_parse(monkeypatch)
|
|
|
|
out = eng._materialize_node_attachments(
|
|
_node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES
|
|
)
|
|
|
|
expected = parse_timeout_for_size(size)
|
|
assert out == [{"id": aid, "mime_type": "text/plain", "content": "PARSED"}]
|
|
assert captured["timeout"] == expected
|
|
# Without the per-call limits the worker self-terminates at the base timeout,
|
|
# so the longer await would buy nothing.
|
|
assert captured["kwargs"] == parse_task_time_limits(expected)
|
|
assert captured["kwargs"]["soft_time_limit"] > settings.DOCUMENT_PARSE_TIMEOUT
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_parse_window_falls_back_to_the_base_timeout_without_a_size(monkeypatch):
|
|
"""A NULL/missing version size keeps the configured base window as the floor."""
|
|
aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf", size=None)
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
captured = _capture_parse(monkeypatch)
|
|
|
|
eng._materialize_node_attachments(
|
|
_node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES
|
|
)
|
|
|
|
assert captured["timeout"] == float(settings.DOCUMENT_PARSE_TIMEOUT)
|
|
assert captured["kwargs"]["soft_time_limit"] == settings.DOCUMENT_PARSE_TIMEOUT
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_preextracted_text_skips_the_parse_worker(monkeypatch):
|
|
"""Text already extracted at upload time is inlined; the parsing worker is never called."""
|
|
aid, rec = _artifact(RUN_ID, "application/pdf", filename="scan.pdf")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
eng = _engine(monkeypatch)
|
|
eng.preextracted_text[aid] = "ALREADY OCRED"
|
|
|
|
import docsgpt.api.user.tasks as tasks
|
|
|
|
monkeypatch.setattr(
|
|
tasks.parse_document,
|
|
"apply_async",
|
|
lambda *a, **k: (_ for _ in ()).throw(AssertionError("must not re-parse")),
|
|
)
|
|
|
|
out = eng._materialize_node_attachments(
|
|
_node_config(input_documents=[aid], file_passing="extract"), "Node", TEXT_TYPES
|
|
)
|
|
|
|
assert out == [{"id": aid, "mime_type": "text/plain", "content": "ALREADY OCRED"}]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _execute_agent_node wiring
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _StubLLM:
|
|
"""Minimal LLM stub exposing the provider's supported-types list (the native source)."""
|
|
|
|
def __init__(self, supported_types):
|
|
self._supported_types = supported_types
|
|
|
|
def get_supported_attachment_types(self):
|
|
return self._supported_types
|
|
|
|
|
|
class _StubNodeAgent:
|
|
"""Captures the factory kwargs and yields a fixed answer (no LLM call)."""
|
|
|
|
def __init__(self, *, supported_types=None, **kwargs):
|
|
self.kwargs = kwargs
|
|
self.tool_executor = None
|
|
self.attachments = []
|
|
self.llm = _StubLLM(supported_types if supported_types is not None else [])
|
|
|
|
def gen(self, _prompt):
|
|
yield {"answer": "ok"}
|
|
|
|
|
|
def _agent_node(input_documents=None, file_passing="auto", node_id="agent_1") -> WorkflowNode:
|
|
config = {
|
|
"agent_type": "classic",
|
|
"system_prompt": "You are a helpful assistant.",
|
|
"prompt_template": "",
|
|
"stream_to_user": False,
|
|
"tools": [],
|
|
}
|
|
if input_documents is not None:
|
|
config["input_documents"] = input_documents
|
|
config["file_passing"] = file_passing
|
|
return WorkflowNode(
|
|
id=node_id,
|
|
workflow_id="wf-1",
|
|
type=NodeType.AGENT,
|
|
title="Agent",
|
|
position={"x": 0, "y": 0},
|
|
config=config,
|
|
)
|
|
|
|
|
|
def _patch_capabilities(monkeypatch):
|
|
"""Stub provider/api-key resolution (capabilities are only fetched for json_schema nodes)."""
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_model_capabilities",
|
|
lambda model_id, user_id=None: None,
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_provider_from_model_id",
|
|
lambda model_id, user_id=None: "openai",
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.core.model_utils.get_api_key_for_provider", lambda _p: "k"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_execute_agent_node_assigns_native_attachment_from_provider_types(monkeypatch):
|
|
"""A node with input_documents assigns a native attachment, decided from the provider's types."""
|
|
aid, rec = _artifact(RUN_ID, "image/png", filename="c.png")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
_patch_capabilities(monkeypatch)
|
|
eng = _engine(monkeypatch)
|
|
|
|
captured: dict = {}
|
|
created: list = []
|
|
|
|
def _create(**kwargs):
|
|
captured.update(kwargs)
|
|
agent = _StubNodeAgent(supported_types=VISION_TYPES, **kwargs)
|
|
created.append(agent)
|
|
return agent
|
|
|
|
monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create))
|
|
|
|
node = _agent_node(input_documents=[aid], file_passing="auto")
|
|
list(eng._execute_agent_node(node))
|
|
|
|
# Attachments are assigned post-construction (BaseAgent reads them at gen time),
|
|
# not passed through factory_kwargs.
|
|
assert "attachments" not in captured
|
|
assert created[0].attachments == [
|
|
{"id": aid, "mime_type": "image/png", "path": rec["versions"][1]["storage_path"]}
|
|
]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_execute_agent_node_native_decision_tracks_provider_types(monkeypatch):
|
|
"""A mime the provider does NOT support falls back to extracted text under ``auto``."""
|
|
aid, rec = _artifact(RUN_ID, "image/png", filename="c.png")
|
|
_patch_repo(monkeypatch, {aid: rec})
|
|
_patch_capabilities(monkeypatch)
|
|
monkeypatch.setattr(
|
|
WorkflowEngine, "_parse_document_text", lambda self, artifact_id, size=None, deadline=None: "EXTRACTED"
|
|
)
|
|
eng = _engine(monkeypatch)
|
|
|
|
created: list = []
|
|
|
|
def _create(**kwargs):
|
|
# Provider reports NO native types, so the registry-agnostic decision
|
|
# must extract rather than send a native-but-empty attachment.
|
|
agent = _StubNodeAgent(supported_types=[], **kwargs)
|
|
created.append(agent)
|
|
return agent
|
|
|
|
monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create))
|
|
|
|
node = _agent_node(input_documents=[aid], file_passing="auto")
|
|
list(eng._execute_agent_node(node))
|
|
|
|
assert created[0].attachments == [
|
|
{"id": aid, "mime_type": "text/plain", "content": "EXTRACTED"}
|
|
]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_execute_agent_node_without_documents_has_no_attachments(monkeypatch):
|
|
"""A node without input_documents leaves attachments empty (no regression)."""
|
|
_patch_capabilities(monkeypatch)
|
|
eng = _engine(monkeypatch)
|
|
|
|
captured: dict = {}
|
|
created: list = []
|
|
|
|
def _create(**kwargs):
|
|
captured.update(kwargs)
|
|
agent = _StubNodeAgent(supported_types=VISION_TYPES, **kwargs)
|
|
created.append(agent)
|
|
return agent
|
|
|
|
monkeypatch.setattr(WorkflowNodeAgentFactory, "create", staticmethod(_create))
|
|
|
|
node = _agent_node(input_documents=None)
|
|
list(eng._execute_agent_node(node))
|
|
|
|
assert "attachments" not in captured
|
|
assert created[0].attachments == []
|
|
|
|
|
|
class _FakeStorage:
|
|
"""Minimal storage stub: get_file(path).read() -> fixed bytes."""
|
|
|
|
def __init__(self, data: bytes) -> None:
|
|
self._data = data
|
|
|
|
def get_file(self, _path):
|
|
return _FakeFile(self._data)
|
|
|
|
|
|
class _FakeFile:
|
|
def __init__(self, data: bytes) -> None:
|
|
self._data = data
|
|
|
|
def read(self) -> bytes:
|
|
return self._data
|