Pass selected run-scoped documents to workflow agent nodes natively

An agent node can list input_documents (short refs like A1, a literal "*" for all
inputs, or state variables holding refs produced by an upstream node). The engine
resolves each through the run-scoped artifact gate and feeds it to the node's LLM
natively when the model supports the format (PDFs included via the provider's
image synthesis), otherwise extracts it to text with Docling. A file_passing
policy (auto/native/extract) and per-node count + per-file byte caps bound cost;
the native decision uses the same supported-types source the provider handler
uses, so a document is never silently dropped. Bytes never enter run state.
This commit is contained in:
Alex committed 2026-06-25 11:35:29 +01:00
1 parent 95c19f0fe7
commit 9fdaea6cd2
5 files changed
+833 -5

No files matched your search

@@ -109,6 +109,84 @@ _EXTRACT_PROGRAM = (
)
def truncate_text_head_tail(text: str, max_bytes: Optional[int] = None) -> str:
"""Bound text to a head+tail byte window so a large file can't flood context."""
cap = int(max_bytes or _MARKDOWN_MAX_BYTES)
if cap <= 0:
return text
encoded = text.encode("utf-8")
if len(encoded) <= cap:
return text
head = cap // 2
tail = cap - head
dropped = len(encoded) - head - tail
head_text = encoded[:head].decode("utf-8", errors="ignore")
tail_text = encoded[-tail:].decode("utf-8", errors="ignore")
return f"{head_text}\n\n...[truncated {dropped} bytes]...\n\n{tail_text}"
def extract_markdown_from_bytes(
data: bytes,
filename: str,
session_id: str,
*,
markdown_max_bytes: Optional[int] = None,
) -> Optional[str]:
"""Run the fixed Docling program in a sandbox session and return the document's markdown, or None.
The bytes ride in as a DATA file the fixed program reads; nothing untrusted is
interpolated into the program. Returns ``None`` when Docling is unavailable or
extraction fails (the caller decides how to degrade).
"""
safe_name = safe_filename(filename) or "document"
token = uuid.uuid4().hex
token_dir = f"extract/{token}"
input_path = f"{token_dir}/inputs/{safe_name}"
params_path = f"{token_dir}/params.json"
result_path = f"{token_dir}/result.json"
params = {
"input_path": input_path,
"markdown_max_bytes": int(markdown_max_bytes or _MARKDOWN_MAX_BYTES),
"max_tables": 0,
}
program = _EXTRACT_PROGRAM.format(params_path=params_path, result_path=result_path)
timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
manager = SandboxCreator.get_manager()
try:
manager.open(session_id, ttl=timeout)
except Exception:
logger.exception("extract_markdown_from_bytes: failed to open sandbox session")
return None
try:
manager.put_file(session_id, input_path, data)
manager.put_file(session_id, params_path, json.dumps(params).encode("utf-8"))
result = manager.exec(session_id, program, timeout=timeout)
if not result.ok:
return None
raw = manager.get_file(session_id, result_path)
except Exception:
logger.exception("extract_markdown_from_bytes: extraction failed")
return None
finally:
manager.remove_path(session_id, token_dir)
try:
manager.close(session_id)
except Exception:
logger.exception("extract_markdown_from_bytes: session close failed")
if not raw:
return None
try:
extracted = json.loads(raw.decode("utf-8"))
except (ValueError, UnicodeDecodeError):
return None
if not isinstance(extracted, dict) or extracted.get("error") or not extracted.get("ok"):
return None
markdown = extracted.get("markdown")
return markdown if isinstance(markdown, str) else None
class DocumentExtractorTool(Tool):
"""Document Extractor
Convert an input document artifact (pdf/docx/pptx/...) to compact, schema-validated structured JSON via Docling.
+9
View File
@@ -49,6 +49,15 @@ class AgentNodeConfig(BaseModel):
retriever: str = ""
model_id: Optional[str] = None
json_schema: Optional[Dict[str, Any]] = None
# Run-scoped documents fed to this node's LLM. Entries are state-var names
# holding artifact refs (single dict or a list of dicts), raw artifact ids,
# short refs (``A1``), or the ``"*"``/``"input_documents"`` token meaning
# "every ref in ``state['input_documents']``".
input_documents: List[str] = Field(default_factory=list)
# How selected documents reach the model: ``auto`` (native when the model
# accepts the mime, else extract to text), ``native`` (force native; raise
# on an unsupported mime), or ``extract`` (always inline extracted text).
file_passing: Literal["auto", "native", "extract"] = "auto"
class CodeNodeConfig(BaseModel):
+182 -5
View File
@@ -250,10 +250,10 @@ class WorkflowEngine:
)
node_api_key = get_api_key_for_provider(node_llm_name) or self.agent.api_key
# Structured output gates on the model's registry capability flags;
# fetch them only when a node json_schema needs the check.
if node_json_schema and node_model_id:
model_capabilities = get_model_capabilities(
node_model_id, user_id=node_user_id
)
model_capabilities = get_model_capabilities(node_model_id, user_id=node_user_id)
if model_capabilities and not model_capabilities.get(
"supports_structured_output", False
):
@@ -294,6 +294,16 @@ class WorkflowEngine:
if getattr(node_agent, "tool_executor", None) is not None:
node_agent.tool_executor.workflow_run_id = self.workflow_run_id
# Decide native-eligibility from the SAME supported-types list the provider
# handler filters on at send time, so a mime is never sent native-but-empty
# and then silently dropped. Read post-construction: BaseAgent consumes
# ``self.attachments`` at gen time, so assigning here takes effect.
node_attachments = self._materialize_node_attachments(
node_config, node.title, self._agent_supported_attachment_types(node_agent)
)
if node_attachments:
node_agent.attachments = node_attachments
full_response_parts: List[str] = []
structured_response_parts: List[str] = []
has_structured_response = False
@@ -477,16 +487,183 @@ class WorkflowEngine:
loaded.append(f"inputs/{filename}")
return loaded
def _materialize_node_attachments(
self,
node_config: AgentNodeConfig,
node_title: str,
supported_types: List[str],
) -> List[Dict[str, Any]]:
"""Resolve a node's selected documents to native/extracted attachment dicts for its LLM."""
from application.agents.tools.artifact_ref import resolve_artifact_id
from application.core.settings import settings
from application.storage.db.repositories.artifacts import ArtifactsRepository
from application.storage.db.session import db_readonly
raw_ids = self._resolve_input_artifact_ids(node_config.input_documents)
if not raw_ids:
return []
supported = set(supported_types)
supports_images = any(t.startswith("image/") for t in supported)
max_files = int(getattr(settings, "WORKFLOW_NODE_NATIVE_MAX_FILES", 5))
max_bytes = int(getattr(settings, "SANDBOX_MAX_INPUT_BYTES", 25 * 1024 * 1024))
# One read-only connection for the whole batch; the resolved-version
# rows are collected, then storage reads happen outside the DB context.
resolved: List[tuple] = []
with db_readonly() as conn:
repo = ArtifactsRepository(conn)
for raw in raw_ids:
# A short ref (A1/...) resolves to an id within this run only; the
# resolved id is re-checked through the run-scoped gate so a forged
# or cross-run ref can never reach another tenant's bytes.
artifact_id = resolve_artifact_id(repo, raw, workflow_run_id=self.workflow_run_id)
artifact = (
repo.get_artifact_in_parent(artifact_id, workflow_run_id=self.workflow_run_id)
if artifact_id is not None
else None
)
if artifact is None:
raise ValueError(
f'Document "{raw}" for node "{node_title}" was not found in this run.'
)
version = repo.get_version(artifact_id, artifact["current_version"])
if not version or not version.get("storage_path"):
raise ValueError(f"input document {artifact_id} has no stored content.")
resolved.append((str(artifact_id), artifact, version))
attachments: List[Dict[str, Any]] = []
native_count = 0
for artifact_id, artifact, version in resolved:
storage_path = version["storage_path"]
mime_type = version.get("mime_type") or "application/octet-stream"
filename = version.get("filename") or artifact.get("title") or artifact_id
size = version.get("size")
if isinstance(size, int) and size > max_bytes:
logger.warning(
"Workflow node %s: document %s (%d bytes) exceeds the %d-byte cap; skipping",
node_title, artifact_id, size, max_bytes,
)
continue
native_ok = self._is_native_mime(mime_type, supported, supports_images)
policy = node_config.file_passing
if policy == "native" and not native_ok:
raise ValueError(
f'Model "{node_config.model_id or self.agent.model_id}" cannot read '
f'"{mime_type}" natively for node "{node_title}".'
)
go_native = policy == "native" or (policy == "auto" and native_ok)
if go_native and native_count >= max_files:
logger.warning(
"Workflow node %s: native file cap (%d) reached; extracting %s instead",
node_title, max_files, filename,
)
go_native = False
if go_native:
# No bytes copied: the provider reads them from storage via ``path``.
attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path})
native_count += 1
else:
content = self._extract_attachment_text(storage_path, mime_type, filename, max_bytes)
if content is None:
logger.warning(
"Workflow node %s: could not extract text from %s; skipping",
node_title, filename,
)
continue
# A non-native mime routes this through ``_append_unsupported_attachments``,
# which inlines ``content`` as text in the system prompt.
attachments.append(
{"id": artifact_id, "mime_type": "text/plain", "content": content}
)
return attachments
@staticmethod
def _agent_supported_attachment_types(node_agent: "BaseAgent") -> List[str]:
"""Return the provider's authoritative supported attachment mime types (handler's source)."""
llm = getattr(node_agent, "llm", None)
getter = getattr(llm, "get_supported_attachment_types", None)
if not callable(getter):
return []
types = getter()
return list(types) if isinstance(types, (list, tuple, set)) else []
@staticmethod
def _is_native_mime(mime_type: str, supported_types: set, supports_images: bool) -> bool:
"""A mime is native if the model accepts it, or it is a PDF a vision model renders to images."""
if mime_type in supported_types:
return True
return mime_type == "application/pdf" and supports_images
def _extract_attachment_text(
self, storage_path: str, mime_type: str, filename: str, max_bytes: int
) -> Optional[str]:
"""Get an attachment's text: inline already-text formats, else extract via Docling; None on failure."""
from application.agents.tools.document_extractor import (
extract_markdown_from_bytes,
truncate_text_head_tail,
)
from application.storage.storage_creator import StorageCreator
try:
data = StorageCreator.get_storage().get_file(storage_path).read()
except Exception:
logger.exception("Workflow node: failed to read document bytes for extraction")
return None
# Defensive size gate: a NULL/missing version size skips the pre-read cap,
# so re-check the actual bytes before inlining or sandbox-extracting.
if len(data) > max_bytes:
logger.warning(
"Workflow node: document at %s (%d bytes) exceeds the %d-byte cap; skipping",
storage_path, len(data), max_bytes,
)
return None
if self._is_inline_text_mime(mime_type):
try:
text = data.decode("utf-8", errors="replace")
except Exception:
return None
# Bound the inlined text to a head+tail window so a large-but-under-cap
# text file can't blow the context (the Docling branch is already bounded).
return truncate_text_head_tail(text)
return extract_markdown_from_bytes(data, filename, self._session_id())
@staticmethod
def _is_inline_text_mime(mime_type: str) -> bool:
"""Already-text formats are inlined directly (no Docling round-trip)."""
if mime_type.startswith("text/"):
return True
return mime_type in ("application/json", "application/xml")
def _resolve_input_artifact_ids(self, inputs: List[str]) -> List[str]:
"""Resolve node ``inputs`` (state-var names holding refs, or raw ids) to artifact ids."""
"""Resolve node ``inputs`` (refs/ids, ``*`` token, or state vars holding a ref or a list of refs)."""
resolved: List[str] = []
for raw in inputs or []:
# ``*`` / ``input_documents`` expands to every run input document.
if isinstance(raw, str) and raw.strip() in ("*", "input_documents"):
resolved.extend(self._input_document_ids())
continue
ref = self.state.get(raw) if isinstance(raw, str) else None
if isinstance(ref, dict) and ref.get("artifact_id"):
resolved.append(str(ref["artifact_id"]))
elif isinstance(ref, list):
# A state var holding a list of ref dicts (e.g. input_documents).
for item in ref:
if isinstance(item, dict) and item.get("artifact_id"):
resolved.append(str(item["artifact_id"]))
elif isinstance(raw, str) and raw.strip():
resolved.append(raw.strip())
return resolved
# Dedup preserving order so ``["*", "A1"]`` / duplicate refs don't attach twice.
return list(dict.fromkeys(resolved))
def _input_document_ids(self) -> List[str]:
"""Return the artifact ids of every ref in ``state['input_documents']``."""
docs = self.state.get("input_documents")
if not isinstance(docs, list):
return []
return [str(d["artifact_id"]) for d in docs if isinstance(d, dict) and d.get("artifact_id")]
def _session_id(self) -> str:
"""Sanitize the run id into the sandbox-gateway charset for the session key."""
+4
View File
@@ -343,6 +343,10 @@ class Settings(BaseSettings):
SANDBOX_MAX_OUTPUT_BYTES: int = 8 * 1024 * 1024 # cap on buffered stdout+stderr per exec
SANDBOX_MAX_FILE_BYTES: int = 10 * 1024 * 1024 # cap on get_file size routed through stdout
SANDBOX_MAX_INPUT_BYTES: int = 25 * 1024 * 1024 # cap on an input document staged into a sandbox session
# Per-agent-node cap on files passed natively to the node's LLM (vision/doc
# inputs). Files past the cap are extracted to text or dropped, not attached
# natively, to bound context/cost. Re-uses SANDBOX_MAX_INPUT_BYTES per file.
WORKFLOW_NODE_NATIVE_MAX_FILES: int = 5
# Runner container resource caps — consumed by the docsgpt-sandbox compose
# service (deployment/sandbox), not by the app client. cgroup CPU/mem caps
# are part of the untrusted-code security boundary.
@@ -0,0 +1,560 @@
"""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 application.agents.workflows.node_agent import WorkflowNodeAgentFactory
from application.agents.workflows.schemas import (
NodeType,
Workflow,
WorkflowGraph,
WorkflowNode,
)
from application.agents.workflows.workflow_engine import WorkflowEngine
from application.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("application.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 application.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(
"application.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(
"application.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_docling(monkeypatch):
"""A non-text mime under extract routes through the Docling helper."""
aid, rec = _artifact(RUN_ID, "application/vnd.openxmlformats", filename="r.docx")
_patch_repo(monkeypatch, {aid: rec})
eng = _engine(monkeypatch)
monkeypatch.setattr(
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _FakeStorage(b"\x00binary")),
)
monkeypatch.setattr(
"application.agents.tools.document_extractor.extract_markdown_from_bytes",
lambda data, filename, session_id, **kw: "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(
"application.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(
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _FakeStorage(b"img-bytes")),
)
# extract of a non-text image falls back to Docling; stub it so it returns text.
monkeypatch.setattr(
"application.agents.tools.document_extractor.extract_markdown_from_bytes",
lambda data, filename, session_id, **kw: "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(
"application.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(
"application.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(
"application.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 application.agents.tools.document_extractor import _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(
"application.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"]}]
# ---------------------------------------------------------------------------
# _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(
"application.core.model_utils.get_model_capabilities",
lambda model_id, user_id=None: None,
)
monkeypatch.setattr(
"application.core.model_utils.get_provider_from_model_id",
lambda model_id, user_id=None: "openai",
)
monkeypatch.setattr(
"application.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(
"application.storage.storage_creator.StorageCreator.get_storage",
staticmethod(lambda: _FakeStorage(b"image-bytes")),
)
monkeypatch.setattr(
"application.agents.tools.document_extractor.extract_markdown_from_bytes",
lambda data, filename, session_id, **kw: "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