mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 20:15:26 +00:00
fix: minor issue fixes
This commit is contained in:
1 parent
99484e02c8
commit
75e07af78f
16 files changed
+950
-43
No files matched your search
@@ -48,6 +48,10 @@ _SESSION_ID_RE = re.compile(r"[^A-Za-z0-9_-]+")
|
||||
# egress-open sandbox code (see ``_json_safe_state``).
|
||||
_CODE_STATE_EXCLUDED_KEYS = frozenset({"chat_history"})
|
||||
|
||||
# Sentinel id stamped on the synthetic text attachment that flags documents a
|
||||
# node skipped past the per-node blocking-extract cap (callers/tests match on it).
|
||||
_EXTRACT_TRUNCATION_ID = "workflow-extract-truncated"
|
||||
|
||||
|
||||
class WorkflowEngine:
|
||||
MAX_EXECUTION_STEPS = 50
|
||||
@@ -511,6 +515,7 @@ class WorkflowEngine:
|
||||
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))
|
||||
extract_max = int(getattr(settings, "WORKFLOW_NODE_EXTRACT_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
|
||||
@@ -539,6 +544,8 @@ class WorkflowEngine:
|
||||
|
||||
attachments: List[Dict[str, Any]] = []
|
||||
native_count = 0
|
||||
extract_count = 0
|
||||
dropped_for_cap = 0
|
||||
for artifact_id, artifact, version in resolved:
|
||||
storage_path = version["storage_path"]
|
||||
mime_type = version.get("mime_type") or "application/octet-stream"
|
||||
@@ -571,6 +578,22 @@ class WorkflowEngine:
|
||||
attachments.append({"id": artifact_id, "mime_type": mime_type, "path": storage_path})
|
||||
native_count += 1
|
||||
else:
|
||||
# Inline-text mimes are read directly (cheap); other mimes route
|
||||
# through the parsing worker -- a blocking, ~120s per-document call.
|
||||
# Cap how many documents a single node sends down that path so a
|
||||
# node referencing many non-native documents (e.g. the ``*`` token)
|
||||
# can't serialize dozens of parses. Inline text is not capped.
|
||||
needs_parse = not self._is_inline_text_mime(mime_type)
|
||||
if needs_parse:
|
||||
# Count (and gate on) the parse ATTEMPT, not the success: a
|
||||
# timed-out/failed parse is the ~120s worst case we must bound,
|
||||
# so it has to consume cap budget too. Otherwise a degraded
|
||||
# parsing backend (every parse fails) never advances the count
|
||||
# and the node keeps issuing blocking parses without limit.
|
||||
if extract_count >= extract_max:
|
||||
dropped_for_cap += 1
|
||||
continue
|
||||
extract_count += 1
|
||||
content = self._extract_attachment_text(
|
||||
artifact_id, storage_path, mime_type, filename, max_bytes
|
||||
)
|
||||
@@ -585,8 +608,26 @@ class WorkflowEngine:
|
||||
attachments.append(
|
||||
{"id": artifact_id, "mime_type": "text/plain", "content": content}
|
||||
)
|
||||
if dropped_for_cap:
|
||||
logger.warning(
|
||||
"Workflow node %s: blocking-extract cap (%d) reached; %d document(s) omitted",
|
||||
node_title, extract_max, dropped_for_cap,
|
||||
)
|
||||
attachments.append(
|
||||
self._extract_truncation_note(node_title, extract_max, dropped_for_cap)
|
||||
)
|
||||
return attachments
|
||||
|
||||
@staticmethod
|
||||
def _extract_truncation_note(node_title: str, cap: int, dropped: int) -> Dict[str, Any]:
|
||||
"""Build a non-fatal text attachment flagging documents skipped past the per-node extract cap."""
|
||||
content = (
|
||||
f'[Notice] Only the first {cap} document(s) for node "{node_title}" were extracted '
|
||||
f"to text; {dropped} additional document(s) were omitted to bound execution time. "
|
||||
"Reference fewer documents, or use a model that reads them natively."
|
||||
)
|
||||
return {"id": _EXTRACT_TRUNCATION_ID, "mime_type": "text/plain", "content": content}
|
||||
|
||||
@staticmethod
|
||||
def _agent_supported_attachment_types(node_agent: "BaseAgent") -> List[str]:
|
||||
"""Return the provider's authoritative supported attachment mime types (handler's source)."""
|
||||
@@ -654,7 +695,13 @@ class WorkflowEngine:
|
||||
args=[artifact_id, {"workflow_run_id": self.workflow_run_id}, user_id, options],
|
||||
queue=queue,
|
||||
)
|
||||
result = async_result.get(timeout=timeout)
|
||||
# A workflow can run inside a Celery worker (scheduled runs / webhooks).
|
||||
# In a prefork worker ``task_join_will_block()`` is process-wide, so the
|
||||
# default ``disable_sync_subtasks=True`` makes ``get()`` raise RuntimeError
|
||||
# ("Never call result.get() within a task!"). The dedicated parsing queue +
|
||||
# separate workers already avoid the real self-deadlock, so opt out
|
||||
# explicitly (mirrors application/agents/tools/read_document.py).
|
||||
result = async_result.get(timeout=timeout, disable_sync_subtasks=False)
|
||||
except (CeleryTimeoutError, TimeoutError):
|
||||
logger.warning("Workflow node: document parse timed out for %s", artifact_id)
|
||||
return None
|
||||
|
||||
@@ -5,7 +5,15 @@ from __future__ import annotations
|
||||
import re
|
||||
from typing import Optional
|
||||
|
||||
from flask import current_app, jsonify, make_response, redirect, request
|
||||
from flask import (
|
||||
Response,
|
||||
current_app,
|
||||
jsonify,
|
||||
make_response,
|
||||
redirect,
|
||||
request,
|
||||
stream_with_context,
|
||||
)
|
||||
from flask_restx import Namespace, Resource
|
||||
|
||||
from application.api import api
|
||||
@@ -183,6 +191,53 @@ class GetArtifact(Resource):
|
||||
current_app.logger.error(f"Error retrieving artifact: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False}), 400)
|
||||
|
||||
@api.doc(description="Delete an artifact and all its versions (owner only)")
|
||||
def delete(self, artifact_id: str):
|
||||
if not looks_like_uuid(artifact_id):
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Artifact not found"}), 404
|
||||
)
|
||||
user_id = resolve_authenticated_user()
|
||||
try:
|
||||
with db_session() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.get_artifact(artifact_id)
|
||||
if artifact is None:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Artifact not found"}), 404
|
||||
)
|
||||
# Delete is a WRITE: only the parent owner may delete; share
|
||||
# links / read-only collaborators are denied (read access only).
|
||||
if not authorize_artifact_write(conn, artifact, user_id):
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "Forbidden"}), 403
|
||||
)
|
||||
storage_paths = repo.delete_artifact(artifact_id)
|
||||
# Reap the bytes best-effort AFTER the row delete commits.
|
||||
_reap_storage(storage_paths)
|
||||
return make_response(jsonify({"success": True}), 200)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error deleting artifact: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False}), 500)
|
||||
|
||||
|
||||
def _reap_storage(paths: list) -> None:
|
||||
"""Best-effort delete artifact bytes after the DB rows are gone; never raises."""
|
||||
if not paths:
|
||||
return
|
||||
try:
|
||||
storage = StorageCreator.get_storage()
|
||||
except Exception:
|
||||
current_app.logger.warning("artifact delete: storage unavailable", exc_info=True)
|
||||
return
|
||||
for path in paths:
|
||||
try:
|
||||
storage.delete_file(path)
|
||||
except Exception:
|
||||
current_app.logger.warning(
|
||||
"artifact delete: failed to delete bytes %s", path, exc_info=True
|
||||
)
|
||||
|
||||
|
||||
@artifacts_ns.route("/artifacts/<artifact_id>/versions/<int:version>")
|
||||
class GetArtifactVersion(Resource):
|
||||
@@ -294,13 +349,26 @@ class DownloadArtifact(Resource):
|
||||
)
|
||||
return redirect(url, code=302)
|
||||
|
||||
# Stream the bytes in chunks instead of buffering the whole object in
|
||||
# worker memory (artifacts can be many MB); close the handle when done.
|
||||
file_obj = storage.get_file(storage_path)
|
||||
response = make_response(file_obj.read())
|
||||
response.headers.set("Content-Type", mime_type)
|
||||
response.headers.set(
|
||||
"Content-Disposition", f'attachment; filename="{filename}"'
|
||||
|
||||
def _stream():
|
||||
try:
|
||||
for chunk in iter(lambda: file_obj.read(65536), b""):
|
||||
yield chunk
|
||||
finally:
|
||||
close = getattr(file_obj, "close", None)
|
||||
if callable(close):
|
||||
close()
|
||||
|
||||
return Response(
|
||||
stream_with_context(_stream()),
|
||||
mimetype=mime_type,
|
||||
headers={
|
||||
"Content-Disposition": f'attachment; filename="{filename}"'
|
||||
},
|
||||
)
|
||||
return response
|
||||
except FileNotFoundError:
|
||||
return make_response(
|
||||
jsonify({"success": False, "message": "File not found"}), 404
|
||||
|
||||
@@ -254,7 +254,29 @@ def process_agent_webhook(self, agent_id, payload, idempotency_key=None):
|
||||
# separate parsing worker and never self-deadlocks the awaiting worker.
|
||||
@celery.task(bind=True, acks_late=False, autoretry_for=())
|
||||
def parse_document(self, artifact_id, parent, user_id, options=None):
|
||||
return parse_document_worker(self, artifact_id, parent, user_id, options or {})
|
||||
"""Parse an input artifact on the parsing queue; self-terminate at the soft time limit."""
|
||||
from celery.exceptions import SoftTimeLimitExceeded
|
||||
|
||||
try:
|
||||
return parse_document_worker(self, artifact_id, parent, user_id, options or {})
|
||||
except SoftTimeLimitExceeded:
|
||||
# A pathological/malicious document must not pin a parsing-worker slot past the
|
||||
# window the caller already abandoned. Return the worker's clean error shape so
|
||||
# the slot frees and the Redis result backend still gets a terminal result.
|
||||
limit = getattr(self, "soft_time_limit", None)
|
||||
suffix = f" after {int(limit)}s" if limit else ""
|
||||
return {"status": "error", "error": f"document parsing timed out{suffix}."}
|
||||
|
||||
|
||||
# Bind the soft limit to DOCUMENT_PARSE_TIMEOUT (the same window read_document awaits) so
|
||||
# the prefork worker self-terminates a runaway parse instead of pinning the slot; the hard
|
||||
# limit is the SIGKILL backstop if the soft handler can't unwind in time.
|
||||
try:
|
||||
from application.core.settings import settings as _parse_settings
|
||||
parse_document.soft_time_limit = int(_parse_settings.DOCUMENT_PARSE_TIMEOUT)
|
||||
parse_document.time_limit = parse_document.soft_time_limit + 30
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@celery.task(**DURABLE_TASK)
|
||||
|
||||
@@ -362,6 +362,12 @@ class Settings(BaseSettings):
|
||||
# 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
|
||||
# Per-agent-node cap on documents extracted to text via the parsing worker.
|
||||
# Each non-native, non-text document issues a separate blocking parse, so a
|
||||
# node referencing many documents (e.g. the ``*`` token) is bounded here to
|
||||
# avoid serializing dozens of parses; documents past the cap are skipped with
|
||||
# a truncation note instead of extracted.
|
||||
WORKFLOW_NODE_EXTRACT_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.
|
||||
|
||||
@@ -1000,6 +1000,10 @@ class LLMHandler(ABC):
|
||||
# They will be added on resume to keep call/result pairs together.
|
||||
continue
|
||||
|
||||
# One assistant(tool_calls) message per call: track whether the
|
||||
# success path already appended it so the except below doesn't
|
||||
# add a second one when create_tool_message fails post-append.
|
||||
assistant_appended = False
|
||||
try:
|
||||
self.tool_calls.append(call)
|
||||
tool_executor_gen = agent._execute_tool_action(tools_dict, call)
|
||||
@@ -1041,41 +1045,53 @@ class LLMHandler(ABC):
|
||||
if reasoning_content:
|
||||
assistant_msg["reasoning_content"] = reasoning_content
|
||||
updated_messages.append(assistant_msg)
|
||||
assistant_appended = True
|
||||
|
||||
updated_messages.append(self.create_tool_message(call, tool_response))
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing tool: {str(e)}", exc_info=True)
|
||||
# The error tool message's tool_call_id must match the assistant
|
||||
# tool_call id that precedes it. When the success path already
|
||||
# appended one (create_tool_message failed afterwards) that id is
|
||||
# the executor-returned ``call_id``; otherwise the except builds
|
||||
# its own assistant message below from ``call.id``.
|
||||
error_id = call_id if assistant_appended else call.id
|
||||
error_call = ToolCall(
|
||||
id=call.id, name=call.name, arguments=call.arguments
|
||||
id=error_id, name=call.name, arguments=call.arguments
|
||||
)
|
||||
error_response = f"Error executing tool: {str(e)}"
|
||||
# Mirror the success path: a role:"tool" message must follow
|
||||
# an assistant message carrying its tool_calls, or the next
|
||||
# provider completion 400s ("'tool' must be a response to a
|
||||
# preceding message with 'tool_calls'").
|
||||
args_str = (
|
||||
json.dumps(call.arguments)
|
||||
if isinstance(call.arguments, dict)
|
||||
else call.arguments
|
||||
)
|
||||
tool_call_obj = {
|
||||
"id": call.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": call.name,
|
||||
"arguments": args_str,
|
||||
},
|
||||
}
|
||||
if call.thought_signature:
|
||||
tool_call_obj["thought_signature"] = call.thought_signature
|
||||
assistant_msg: Dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [tool_call_obj],
|
||||
}
|
||||
if reasoning_content:
|
||||
assistant_msg["reasoning_content"] = reasoning_content
|
||||
updated_messages.append(assistant_msg)
|
||||
# preceding message with 'tool_calls'"). Skip it when the
|
||||
# success path already appended one for this call — a
|
||||
# create_tool_message failure after that append would
|
||||
# otherwise duplicate the assistant message and 400 the
|
||||
# same way an orphan tool message does.
|
||||
if not assistant_appended:
|
||||
args_str = (
|
||||
json.dumps(call.arguments)
|
||||
if isinstance(call.arguments, dict)
|
||||
else call.arguments
|
||||
)
|
||||
tool_call_obj = {
|
||||
"id": call.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": call.name,
|
||||
"arguments": args_str,
|
||||
},
|
||||
}
|
||||
if call.thought_signature:
|
||||
tool_call_obj["thought_signature"] = call.thought_signature
|
||||
assistant_msg: Dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [tool_call_obj],
|
||||
}
|
||||
if reasoning_content:
|
||||
assistant_msg["reasoning_content"] = reasoning_content
|
||||
updated_messages.append(assistant_msg)
|
||||
|
||||
error_message = self.create_tool_message(error_call, error_response)
|
||||
updated_messages.append(error_message)
|
||||
|
||||
@@ -139,6 +139,17 @@ def _parse_to_text(parser: Any, path: Path) -> str:
|
||||
return str(parsed)
|
||||
|
||||
|
||||
def _is_docling_parser(parser: Any) -> bool:
|
||||
"""True when ``parser`` is Docling-backed (collecting tables would otherwise re-convert)."""
|
||||
if parser is None:
|
||||
return False
|
||||
try:
|
||||
from application.parser.file.docling_parser import DoclingParser
|
||||
except Exception:
|
||||
return False
|
||||
return isinstance(parser, DoclingParser)
|
||||
|
||||
|
||||
def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Bound a single table's rows and cell sizes so one giant table can't bloat context."""
|
||||
|
||||
@@ -159,13 +170,27 @@ def _compact_table(table: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return compact
|
||||
|
||||
|
||||
def _docling_structured(path: Path, *, ocr_enabled: bool, include_tables: bool) -> Dict[str, Any]:
|
||||
"""Convert a document with Docling and return markdown + structured dict + bounded tables."""
|
||||
from docling.document_converter import DocumentConverter
|
||||
def _docling_structured(
|
||||
path: Path, *, ocr_enabled: bool, include_tables: bool, parser: Any = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Convert a document with Docling and return markdown + structured dict + bounded tables.
|
||||
|
||||
converter = DocumentConverter()
|
||||
doc = converter.convert(str(path)).document
|
||||
markdown = doc.export_to_markdown()
|
||||
When ``parser`` is the configured ``DoclingParser`` (the collapse-the-double-convert
|
||||
path), reuse ITS converter + export so OCR/pipeline options and the export fallback
|
||||
are honored and the content matches the legacy single-parse output exactly; otherwise
|
||||
fall back to a vanilla converter.
|
||||
"""
|
||||
if _is_docling_parser(parser):
|
||||
if getattr(parser, "_converter", None) is None:
|
||||
parser._init_parser()
|
||||
doc = parser._converter.convert(str(path)).document
|
||||
markdown = parser._export_content(doc)
|
||||
else:
|
||||
from docling.document_converter import DocumentConverter
|
||||
|
||||
converter = DocumentConverter()
|
||||
doc = converter.convert(str(path)).document
|
||||
markdown = doc.export_to_markdown()
|
||||
structured = doc.export_to_dict()
|
||||
tables: List[Dict[str, Any]] = []
|
||||
if include_tables:
|
||||
@@ -323,6 +348,28 @@ def _shape(
|
||||
}
|
||||
|
||||
parser = _pick_parser(suffix, ocr_enabled=ocr_enabled, engine=engine)
|
||||
wants_tables = include_tables and engine != "fast" and output != "chunks"
|
||||
|
||||
# A Docling-backed parser already converts the whole document to produce its text.
|
||||
# When tables are also requested, reuse that single conversion for both the markdown
|
||||
# content and the tables instead of converting a second time just to collect tables
|
||||
# (Docling/torch conversion dominates the cost, so a re-convert ~doubles latency).
|
||||
if wants_tables and _is_docling_parser(parser):
|
||||
try:
|
||||
extracted = _docling_structured(
|
||||
path, ocr_enabled=ocr_enabled, include_tables=True, parser=parser
|
||||
)
|
||||
text = extracted["markdown"]
|
||||
tables: List[Dict[str, Any]] = [_compact_table(t) for t in extracted["tables"]]
|
||||
except Exception:
|
||||
text, tables = _parse_to_text(parser, path), []
|
||||
text = _apply_pages(text, pages)
|
||||
bounded, truncated = _bounded(text, max_chars)
|
||||
payload: Dict[str, Any] = {"output": output, "content": bounded, "truncated": truncated}
|
||||
if tables:
|
||||
payload["tables"] = tables
|
||||
return payload
|
||||
|
||||
if parser is None:
|
||||
# A whitelisted extension with no dedicated parser (e.g. .txt) reads as plain
|
||||
# text, matching SimpleDirectoryReader's standard-read fallback.
|
||||
@@ -335,7 +382,7 @@ def _shape(
|
||||
return {"output": "chunks", "chunks": _to_chunks(text, max_chars), "truncated": False}
|
||||
|
||||
tables: List[Dict[str, Any]] = []
|
||||
if include_tables and engine != "fast":
|
||||
if wants_tables:
|
||||
try:
|
||||
tables = [_compact_table(t) for t in _docling_structured(
|
||||
path, ocr_enabled=ocr_enabled, include_tables=True)["tables"]]
|
||||
|
||||
@@ -301,6 +301,72 @@ class ArtifactsRepository:
|
||||
).fetchone()
|
||||
return int(row[0]) if row is not None else 0
|
||||
|
||||
def delete_artifact(self, artifact_id: str) -> list[str]:
|
||||
"""Delete an artifact (+ its versions, via FK cascade); return its stored paths.
|
||||
|
||||
The caller deletes the returned storage objects best-effort AFTER commit;
|
||||
the DB delete also frees the owner's count/byte quota.
|
||||
"""
|
||||
paths = [
|
||||
v["storage_path"]
|
||||
for v in self.list_versions(artifact_id)
|
||||
if v.get("storage_path")
|
||||
]
|
||||
self._conn.execute(
|
||||
text("DELETE FROM artifacts WHERE id = CAST(:id AS uuid)"),
|
||||
{"id": artifact_id},
|
||||
)
|
||||
return paths
|
||||
|
||||
def storage_paths_for_conversation(self, conversation_id: str) -> list[str]:
|
||||
"""Return every stored version path for a conversation's artifacts (for byte cleanup)."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT v.storage_path FROM artifact_versions v "
|
||||
"JOIN artifacts a ON a.id = v.artifact_id "
|
||||
"WHERE a.conversation_id = CAST(:cid AS uuid) "
|
||||
"AND v.storage_path IS NOT NULL"
|
||||
),
|
||||
{"cid": conversation_id},
|
||||
)
|
||||
return [r[0] for r in result.fetchall()]
|
||||
|
||||
def storage_paths_for_user_conversations(self, user_id: str) -> list[str]:
|
||||
"""Return every stored version path for artifacts under a user's conversations."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT v.storage_path FROM artifact_versions v "
|
||||
"JOIN artifacts a ON a.id = v.artifact_id "
|
||||
"JOIN conversations c ON c.id = a.conversation_id "
|
||||
"WHERE c.user_id = :user_id AND v.storage_path IS NOT NULL"
|
||||
),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
return [r[0] for r in result.fetchall()]
|
||||
|
||||
def delete_for_conversation(self, conversation_id: str) -> int:
|
||||
"""Delete all artifacts (+ versions, via cascade) parented to a conversation.
|
||||
|
||||
``conversation_id`` is a bare uuid column (no FK), so the parent delete
|
||||
does not cascade — callers invoke this to reclaim the rows + their quota.
|
||||
"""
|
||||
result = self._conn.execute(
|
||||
text("DELETE FROM artifacts WHERE conversation_id = CAST(:cid AS uuid)"),
|
||||
{"cid": conversation_id},
|
||||
)
|
||||
return result.rowcount
|
||||
|
||||
def delete_for_user_conversations(self, user_id: str) -> int:
|
||||
"""Delete every artifact parented to one of a user's conversations (+ versions)."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"DELETE FROM artifacts WHERE conversation_id IN "
|
||||
"(SELECT id FROM conversations WHERE user_id = :user_id)"
|
||||
),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
return result.rowcount
|
||||
|
||||
def get_version(self, artifact_id: str, version: int) -> Optional[dict]:
|
||||
"""Fetch a single version row by artifact id and version number."""
|
||||
result = self._conn.execute(
|
||||
|
||||
@@ -15,6 +15,7 @@ Covers every operation the legacy Mongo code performs on
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
@@ -25,6 +26,8 @@ from application.storage.db.base_repository import looks_like_uuid, row_to_dict
|
||||
from application.storage.db.models import conversations_table, conversation_messages_table
|
||||
from application.storage.db.serialization import PGNativeJSONEncoder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MessageUpdateOutcome(str, Enum):
|
||||
"""Discriminated result of ``update_message_by_id``.
|
||||
@@ -535,6 +538,14 @@ class ConversationsRepository:
|
||||
# a non-UUID id reaches this code path.
|
||||
if not looks_like_uuid(conversation_id):
|
||||
return False
|
||||
# Artifacts are parent-derived: deleting a conversation must take its
|
||||
# artifacts with it (rows + bytes + the quota they consume). conversation_id
|
||||
# is a bare uuid (no FK cascade), so delete the rows explicitly and reap the
|
||||
# stored bytes best-effort.
|
||||
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
||||
|
||||
artifacts = ArtifactsRepository(self._conn)
|
||||
paths = artifacts.storage_paths_for_conversation(conversation_id)
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"DELETE FROM conversations "
|
||||
@@ -542,15 +553,45 @@ class ConversationsRepository:
|
||||
),
|
||||
{"id": conversation_id, "user_id": user_id},
|
||||
)
|
||||
return result.rowcount > 0
|
||||
deleted = result.rowcount > 0
|
||||
if deleted:
|
||||
artifacts.delete_for_conversation(conversation_id)
|
||||
self._reap_artifact_storage(paths)
|
||||
return deleted
|
||||
|
||||
def delete_all_for_user(self, user_id: str) -> int:
|
||||
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
||||
|
||||
# Reap artifacts BEFORE the conversations: delete_for_user_conversations
|
||||
# resolves them via a subquery over the still-present conversation rows.
|
||||
artifacts = ArtifactsRepository(self._conn)
|
||||
paths = artifacts.storage_paths_for_user_conversations(user_id)
|
||||
artifacts.delete_for_user_conversations(user_id)
|
||||
result = self._conn.execute(
|
||||
text("DELETE FROM conversations WHERE user_id = :user_id"),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
self._reap_artifact_storage(paths)
|
||||
return result.rowcount
|
||||
|
||||
@staticmethod
|
||||
def _reap_artifact_storage(paths: list) -> None:
|
||||
"""Best-effort delete artifact bytes whose rows were cascade-removed; never raises."""
|
||||
if not paths:
|
||||
return
|
||||
try:
|
||||
from application.storage.storage_creator import StorageCreator
|
||||
|
||||
storage = StorageCreator.get_storage()
|
||||
except Exception:
|
||||
logger.warning("conversation delete: storage unavailable for artifact cleanup", exc_info=True)
|
||||
return
|
||||
for path in paths:
|
||||
try:
|
||||
storage.delete_file(path)
|
||||
except Exception:
|
||||
logger.warning("conversation delete: failed to delete artifact bytes %s", path, exc_info=True)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Messages
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -18,8 +18,17 @@ services:
|
||||
- CELERY_BROKER_URL=redis://redis:6379/0
|
||||
- CELERY_RESULT_BACKEND=redis://redis:6379/1
|
||||
- POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt
|
||||
# Code-execution runner reached over HTTP + WebSocket (no docker socket).
|
||||
- SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888
|
||||
# Env-scrubbing kernelspec the docsgpt-sandbox image installs.
|
||||
- SANDBOX_KERNEL_NAME=docsgpt-python
|
||||
ports:
|
||||
- "7091:7091"
|
||||
networks:
|
||||
# `default` reaches the sandbox gateway; `data-net` reaches redis/postgres
|
||||
# (which are NOT on the sandbox's networks).
|
||||
- default
|
||||
- data-net
|
||||
volumes:
|
||||
- ../application/indexes:/app/application/indexes
|
||||
- ../application/inputs:/app/application/inputs
|
||||
@@ -42,19 +51,54 @@ services:
|
||||
- CELERY_RESULT_BACKEND=redis://redis:6379/1
|
||||
- API_URL=http://backend:7091
|
||||
- POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt
|
||||
- SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888
|
||||
- SANDBOX_KERNEL_NAME=docsgpt-python
|
||||
networks:
|
||||
- default
|
||||
- data-net
|
||||
depends_on:
|
||||
redis:
|
||||
condition: service_started
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
|
||||
# Always-on code-execution runner (Jupyter Kernel Gateway). In-process kernels,
|
||||
# no Docker socket, no published host port. On the internal-only sandbox-net
|
||||
# plus `default` (for internet egress); deliberately NOT on data-net, so
|
||||
# sandboxed code cannot reach the unauthenticated redis broker (Celery task
|
||||
# injection -> RCE) or postgres. Do NOT add `env_file: ../.env` here -- the
|
||||
# runner needs no app secrets. For untrusted multi-tenant workloads use a
|
||||
# per-session VM via SANDBOX_BACKEND=daytona instead.
|
||||
docsgpt-sandbox:
|
||||
build: ./sandbox
|
||||
mem_limit: ${SANDBOX_MEMORY:-1g}
|
||||
cpus: ${SANDBOX_CPUS:-1.0}
|
||||
pids_limit: 256
|
||||
read_only: true
|
||||
environment:
|
||||
# Keep Jupyter's runtime/connection files on the writable tmpfs.
|
||||
- JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime
|
||||
- JUPYTER_DATA_DIR=/tmp/jupyter-data
|
||||
tmpfs:
|
||||
- /tmp
|
||||
networks:
|
||||
- sandbox-net
|
||||
- default
|
||||
|
||||
redis:
|
||||
image: redis:6-alpine
|
||||
# Data plane only: kept off the sandbox's networks (a reachable, unauthed
|
||||
# broker = Celery task injection -> worker RCE).
|
||||
networks:
|
||||
- data-net
|
||||
ports:
|
||||
- 6379:6379
|
||||
|
||||
postgres:
|
||||
image: postgres:16-alpine
|
||||
# Data plane only (see redis): kept off the sandbox's networks.
|
||||
networks:
|
||||
- data-net
|
||||
environment:
|
||||
- POSTGRES_USER=docsgpt
|
||||
- POSTGRES_PASSWORD=docsgpt
|
||||
@@ -69,5 +113,15 @@ services:
|
||||
timeout: 5s
|
||||
retries: 10
|
||||
|
||||
networks:
|
||||
# Internal-only runner network (no external egress via this net; the sandbox
|
||||
# still reaches the internet via the default bridge).
|
||||
sandbox-net:
|
||||
internal: true
|
||||
# Internal data plane: redis/postgres + backend/worker. The sandbox is NOT
|
||||
# attached, so sandboxed code cannot reach the broker or database.
|
||||
data-net:
|
||||
internal: true
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
@@ -23,8 +23,17 @@ services:
|
||||
- CELERY_RESULT_BACKEND=redis://redis:6379/1
|
||||
- CACHE_REDIS_URL=redis://redis:6379/2
|
||||
- POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt
|
||||
# Code-execution runner reached over HTTP + WebSocket (no docker socket).
|
||||
- SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888
|
||||
# Env-scrubbing kernelspec the docsgpt-sandbox image installs.
|
||||
- SANDBOX_KERNEL_NAME=docsgpt-python
|
||||
ports:
|
||||
- "7091:7091"
|
||||
networks:
|
||||
# `default` reaches the sandbox gateway; `data-net` reaches redis/postgres
|
||||
# (which are NOT on the sandbox's networks).
|
||||
- default
|
||||
- data-net
|
||||
volumes:
|
||||
- ../application/indexes:/app/indexes
|
||||
- ../application/inputs:/app/inputs
|
||||
@@ -49,6 +58,11 @@ services:
|
||||
- API_URL=http://backend:7091
|
||||
- CACHE_REDIS_URL=redis://redis:6379/2
|
||||
- POSTGRES_URI=postgresql://docsgpt:docsgpt@postgres:5432/docsgpt
|
||||
- SANDBOX_GATEWAY_URL=http://docsgpt-sandbox:8888
|
||||
- SANDBOX_KERNEL_NAME=docsgpt-python
|
||||
networks:
|
||||
- default
|
||||
- data-net
|
||||
volumes:
|
||||
- ../application/indexes:/app/indexes
|
||||
- ../application/inputs:/app/inputs
|
||||
@@ -59,13 +73,46 @@ services:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
|
||||
# Always-on code-execution runner (Jupyter Kernel Gateway). In-process kernels,
|
||||
# no Docker socket, no published host port. On the internal-only sandbox-net
|
||||
# plus `default` (for internet egress); deliberately NOT on data-net, so
|
||||
# sandboxed code cannot reach the unauthenticated redis broker (Celery task
|
||||
# injection -> RCE) or postgres. Do NOT add `env_file: ../.env` here -- the
|
||||
# runner needs no app secrets. For untrusted multi-tenant workloads use a
|
||||
# per-session VM via SANDBOX_BACKEND=daytona instead.
|
||||
docsgpt-sandbox:
|
||||
# Built from the repo (no published arc53/docsgpt-sandbox image exists yet); the
|
||||
# hub compose ships alongside ./sandbox, so this stays self-contained and the
|
||||
# stack always starts. Switch to a published image once CI builds/pushes one.
|
||||
build: ./sandbox
|
||||
mem_limit: ${SANDBOX_MEMORY:-1g}
|
||||
cpus: ${SANDBOX_CPUS:-1.0}
|
||||
pids_limit: 256
|
||||
read_only: true
|
||||
environment:
|
||||
# Keep Jupyter's runtime/connection files on the writable tmpfs.
|
||||
- JUPYTER_RUNTIME_DIR=/tmp/jupyter-runtime
|
||||
- JUPYTER_DATA_DIR=/tmp/jupyter-data
|
||||
tmpfs:
|
||||
- /tmp
|
||||
networks:
|
||||
- sandbox-net
|
||||
- default
|
||||
|
||||
redis:
|
||||
image: redis:6-alpine
|
||||
# Data plane only: kept off the sandbox's networks (a reachable, unauthed
|
||||
# broker = Celery task injection -> worker RCE).
|
||||
networks:
|
||||
- data-net
|
||||
ports:
|
||||
- 6379:6379
|
||||
|
||||
postgres:
|
||||
image: postgres:16-alpine
|
||||
# Data plane only (see redis): kept off the sandbox's networks.
|
||||
networks:
|
||||
- data-net
|
||||
environment:
|
||||
- POSTGRES_USER=docsgpt
|
||||
- POSTGRES_PASSWORD=docsgpt
|
||||
@@ -80,5 +127,15 @@ services:
|
||||
timeout: 5s
|
||||
retries: 10
|
||||
|
||||
networks:
|
||||
# Internal-only runner network (no external egress via this net; the sandbox
|
||||
# still reaches the internet via the default bridge).
|
||||
sandbox-net:
|
||||
internal: true
|
||||
# Internal data plane: redis/postgres + backend/worker. The sandbox is NOT
|
||||
# attached, so sandboxed code cannot reach the broker or database.
|
||||
data-net:
|
||||
internal: true
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
@@ -16,7 +16,11 @@ import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.agents.workflow_agent import WorkflowAgent, _MAX_INPUT_DOCUMENTS
|
||||
from application.agents.workflows.workflow_engine import WorkflowEngine
|
||||
from application.agents.workflows.schemas import AgentNodeConfig
|
||||
from application.agents.workflows.workflow_engine import (
|
||||
_EXTRACT_TRUNCATION_ID,
|
||||
WorkflowEngine,
|
||||
)
|
||||
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
||||
from application.storage.db.repositories.workflow_runs import WorkflowRunsRepository
|
||||
from application.storage.local import LocalStorage
|
||||
@@ -350,3 +354,142 @@ def test_no_attachments_run_still_works(pg_engine, tmp_path, monkeypatch):
|
||||
).scalar()
|
||||
assert run is not None
|
||||
assert n == 0
|
||||
|
||||
|
||||
def test_extract_parse_opts_out_of_sync_subtask_guard(monkeypatch):
|
||||
"""_parse_document_text awaits with disable_sync_subtasks=False so it works inside a Celery worker."""
|
||||
agent = _agent(str(uuid.uuid4()), [])
|
||||
engine = WorkflowEngine.__new__(WorkflowEngine)
|
||||
engine.agent = agent
|
||||
engine.workflow_run_id = "run-extract"
|
||||
|
||||
import application.api.user.tasks as tasks
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
class _FakeAsyncResult:
|
||||
def __init__(self):
|
||||
self.get_kwargs = None
|
||||
|
||||
def get(self, timeout=None, disable_sync_subtasks=True):
|
||||
self.get_kwargs = {"timeout": timeout, "disable_sync_subtasks": disable_sync_subtasks}
|
||||
return {"status": "ok", "content": "parsed markdown"}
|
||||
|
||||
def _apply_async(args=None, queue=None, **kw):
|
||||
result = _FakeAsyncResult()
|
||||
captured["result"] = result
|
||||
captured["args"] = args
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
|
||||
|
||||
out = engine._parse_document_text("artifact-xyz")
|
||||
|
||||
assert out == "parsed markdown"
|
||||
# A prefork worker's task_join_will_block() is process-wide, so the await must
|
||||
# opt out of the guard or get() raises RuntimeError("Never call result.get()...").
|
||||
assert captured["result"].get_kwargs["disable_sync_subtasks"] is False
|
||||
# The run-scoped parent + resolved id reached the parsing task.
|
||||
assert captured["args"][0] == "artifact-xyz"
|
||||
assert captured["args"][1] == {"workflow_run_id": "run-extract"}
|
||||
|
||||
|
||||
def test_node_extract_path_capped_with_truncation_note(pg_engine, tmp_path, monkeypatch):
|
||||
"""A node referencing more docs than the extract cap parses only up to the cap and notes the rest."""
|
||||
storage = _wire(pg_engine, tmp_path, monkeypatch)
|
||||
wf_id = _make_workflow(pg_engine)
|
||||
|
||||
# docx is non-native (no vision) and not inline-text, so each routes through
|
||||
# the blocking parsing worker -- the path the per-node cap must bound.
|
||||
docx = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
attachments = [
|
||||
_stage_attachment(storage, f"doc-{i}".encode(), f"f{i}.docx", docx)
|
||||
for i in range(4)
|
||||
]
|
||||
agent = _agent(wf_id, attachments)
|
||||
_patch_engine(monkeypatch)
|
||||
list(agent._gen_inner("summarize", log_context=None))
|
||||
engine = _RecordingEngine.instances[-1]
|
||||
|
||||
# Cap the blocking-extract path below the doc count so the overflow truncates.
|
||||
from application.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 2, raising=False)
|
||||
|
||||
# Stub the parsing worker so each non-text doc "parses" without a broker, and
|
||||
# count the blocking calls to prove the overflow docs are never enqueued.
|
||||
import application.api.user.tasks as tasks
|
||||
|
||||
parse_calls = {"n": 0}
|
||||
|
||||
class _R:
|
||||
def get(self, timeout=None, disable_sync_subtasks=True):
|
||||
return {"status": "ok", "content": "PARSED"}
|
||||
|
||||
def _apply_async(args=None, queue=None, **kw):
|
||||
parse_calls["n"] += 1
|
||||
return _R()
|
||||
|
||||
monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
|
||||
|
||||
node_config = AgentNodeConfig(input_documents=["*"])
|
||||
out = engine._materialize_node_attachments(node_config, "Reviewer", supported_types=[])
|
||||
|
||||
notes = [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID]
|
||||
extracted = [a for a in out if a.get("id") != _EXTRACT_TRUNCATION_ID]
|
||||
# Only the cap was extracted; the remaining docs were never sent to the worker.
|
||||
assert len(extracted) == 2
|
||||
assert parse_calls["n"] == 2
|
||||
# A single non-fatal truncation note is appended to the node's inlined text.
|
||||
assert len(notes) == 1
|
||||
assert notes[0]["mime_type"] == "text/plain"
|
||||
assert "omitted" in notes[0]["content"].lower()
|
||||
|
||||
|
||||
def test_node_extract_cap_bounds_parse_attempts_even_when_every_parse_times_out(
|
||||
pg_engine, tmp_path, monkeypatch
|
||||
):
|
||||
"""The cap must bound parse ATTEMPTS, not successes: a degraded backend where every
|
||||
parse times out (~120s each) must still issue at most the cap's worth of blocking calls."""
|
||||
from celery.exceptions import TimeoutError as CeleryTimeoutError
|
||||
|
||||
storage = _wire(pg_engine, tmp_path, monkeypatch)
|
||||
wf_id = _make_workflow(pg_engine)
|
||||
|
||||
docx = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
attachments = [
|
||||
_stage_attachment(storage, f"doc-{i}".encode(), f"f{i}.docx", docx)
|
||||
for i in range(4)
|
||||
]
|
||||
agent = _agent(wf_id, attachments)
|
||||
_patch_engine(monkeypatch)
|
||||
list(agent._gen_inner("summarize", log_context=None))
|
||||
engine = _RecordingEngine.instances[-1]
|
||||
|
||||
from application.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "WORKFLOW_NODE_EXTRACT_MAX_FILES", 2, raising=False)
|
||||
|
||||
import application.api.user.tasks as tasks
|
||||
|
||||
parse_calls = {"n": 0}
|
||||
|
||||
class _R:
|
||||
def get(self, timeout=None, disable_sync_subtasks=True):
|
||||
raise CeleryTimeoutError() # the ~120s worst case the cap must bound
|
||||
|
||||
def _apply_async(args=None, queue=None, **kw):
|
||||
parse_calls["n"] += 1
|
||||
return _R()
|
||||
|
||||
monkeypatch.setattr(tasks.parse_document, "apply_async", _apply_async)
|
||||
|
||||
node_config = AgentNodeConfig(input_documents=["*"])
|
||||
out = engine._materialize_node_attachments(node_config, "Reviewer", supported_types=[])
|
||||
|
||||
# Every parse timed out (nothing extracted), but blocking attempts were bounded.
|
||||
extracted = [a for a in out if a.get("id") != _EXTRACT_TRUNCATION_ID]
|
||||
assert extracted == []
|
||||
assert parse_calls["n"] == 2 # not 4 -- failed parses still consume cap budget
|
||||
notes = [a for a in out if a.get("id") == _EXTRACT_TRUNCATION_ID]
|
||||
assert len(notes) == 1
|
||||
@@ -297,8 +297,9 @@ class TestSharedAccess:
|
||||
# Download
|
||||
# ---------------------------------------------------------------------------
|
||||
class _FakeStorage:
|
||||
def __init__(self, data: bytes):
|
||||
def __init__(self, data: bytes = b""):
|
||||
self._data = data
|
||||
self.deleted: list = []
|
||||
|
||||
def get_file(self, path):
|
||||
return io.BytesIO(self._data)
|
||||
@@ -306,6 +307,10 @@ class _FakeStorage:
|
||||
def generate_presigned_url(self, path, expires_in=300):
|
||||
return f"https://signed.example/{path}?exp={expires_in}"
|
||||
|
||||
def delete_file(self, path):
|
||||
self.deleted.append(path)
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDownloadArtifact:
|
||||
@@ -343,6 +348,27 @@ class TestDownloadArtifact:
|
||||
assert resp.headers["Content-Disposition"] == 'attachment; filename="deck.pptx"'
|
||||
assert resp.headers["Content-Type"] == "application/vnd.ms-powerpoint"
|
||||
|
||||
def test_local_download_is_streamed_not_buffered(
|
||||
self, _patch_db, flask_app, token_owner, monkeypatch
|
||||
):
|
||||
# The response body must be a stream (generator), not the whole object
|
||||
# buffered into memory via make_response(file_obj.read()).
|
||||
from application.api.user.artifacts.routes import DownloadArtifact
|
||||
|
||||
art = self._seed(_patch_db)
|
||||
storage = _FakeStorage(b"Z" * 200_000)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.artifacts.routes.StorageCreator.get_storage",
|
||||
lambda: storage,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.artifacts.routes.settings.URL_STRATEGY",
|
||||
"backend", raising=False,
|
||||
)
|
||||
resp = _call(flask_app, DownloadArtifact, art["id"], token=token_owner)
|
||||
assert resp.is_streamed
|
||||
assert resp.get_data() == b"Z" * 200_000
|
||||
|
||||
def test_s3_strategy_redirects_to_presigned(
|
||||
self, _patch_db, flask_app, token_owner, monkeypatch
|
||||
):
|
||||
@@ -506,6 +532,91 @@ class TestRestoreArtifact:
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestDeleteArtifact:
|
||||
def _seed(self, conn, **over):
|
||||
conv = _make_conversation(conn)
|
||||
defaults = dict(
|
||||
conversation_id=str(conv["id"]),
|
||||
filename="f.bin",
|
||||
storage_path="inputs/owner/artifacts/x/v1/f.bin",
|
||||
)
|
||||
defaults.update(over)
|
||||
return _make_artifact(conn, **defaults)
|
||||
|
||||
def test_owner_deletes_and_reaps_bytes(
|
||||
self, _patch_db, flask_app, token_owner, monkeypatch
|
||||
):
|
||||
from application.api.user.artifacts.routes import GetArtifact
|
||||
|
||||
art = self._seed(_patch_db)
|
||||
storage = _FakeStorage()
|
||||
monkeypatch.setattr(
|
||||
"application.api.user.artifacts.routes.StorageCreator.get_storage",
|
||||
lambda: storage,
|
||||
)
|
||||
resp = _call(flask_app, GetArtifact, art["id"], token=token_owner, method="delete")
|
||||
assert resp.status_code == 200
|
||||
assert ArtifactsRepository(_patch_db).get_artifact(art["id"]) is None
|
||||
assert "inputs/owner/artifacts/x/v1/f.bin" in storage.deleted
|
||||
|
||||
def test_stranger_denied(self, _patch_db, flask_app):
|
||||
from application.api.user.artifacts.routes import GetArtifact
|
||||
|
||||
art = self._seed(_patch_db)
|
||||
resp = _call(
|
||||
flask_app, GetArtifact, art["id"], token={"sub": STRANGER}, method="delete"
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
assert ArtifactsRepository(_patch_db).get_artifact(art["id"]) is not None
|
||||
|
||||
def test_share_token_holder_denied(self, _patch_db, flask_app):
|
||||
from application.api.user.artifacts.routes import GetArtifact
|
||||
|
||||
conv = _make_conversation(_patch_db)
|
||||
art = _make_artifact(_patch_db, conversation_id=str(conv["id"]))
|
||||
share = SharedConversationsRepository(_patch_db).create(str(conv["id"]), OWNER)
|
||||
resp = _call(
|
||||
flask_app, GetArtifact, art["id"], token=None,
|
||||
query={"share_token": str(share["uuid"])}, method="delete",
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
assert ArtifactsRepository(_patch_db).get_artifact(art["id"]) is not None
|
||||
|
||||
def test_unknown_artifact_404(self, _patch_db, flask_app, token_owner):
|
||||
from application.api.user.artifacts.routes import GetArtifact
|
||||
|
||||
resp = _call(
|
||||
flask_app, GetArtifact, str(uuid.uuid4()), token=token_owner, method="delete"
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestConversationDeleteReapsArtifacts:
|
||||
def test_delete_conversation_removes_artifacts_frees_quota_and_reaps_bytes(
|
||||
self, _patch_db, flask_app, monkeypatch
|
||||
):
|
||||
conv = _make_conversation(_patch_db)
|
||||
art = _make_artifact(
|
||||
_patch_db, conversation_id=str(conv["id"]),
|
||||
storage_path="k/reap.bin", filename="r.bin", size=42,
|
||||
)
|
||||
storage = _FakeStorage()
|
||||
monkeypatch.setattr(
|
||||
"application.storage.storage_creator.StorageCreator.get_storage",
|
||||
lambda: storage,
|
||||
)
|
||||
|
||||
deleted = ConversationsRepository(_patch_db).delete(str(conv["id"]), OWNER)
|
||||
assert deleted is True
|
||||
|
||||
repo = ArtifactsRepository(_patch_db)
|
||||
assert repo.get_artifact(art["id"]) is None # cascade-removed with the parent
|
||||
assert repo.count_for_user(OWNER) == 0 # quota freed
|
||||
assert "k/reap.bin" in storage.deleted # bytes reaped
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Malformed artifact_id path segment -> 404 (no CAST(:id AS uuid) txn poison)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -327,6 +327,53 @@ class TestMcpOauthTask:
|
||||
assert result == {"url": "http://auth"}
|
||||
|
||||
|
||||
class TestParseDocumentTask:
|
||||
"""parse_document runs on the parsing queue under a bounded time limit."""
|
||||
|
||||
@pytest.mark.unit
|
||||
@patch("application.api.user.tasks.parse_document_worker")
|
||||
def test_calls_parse_document_worker(self, mock_worker):
|
||||
from application.api.user.tasks import parse_document
|
||||
|
||||
mock_worker.return_value = {"status": "ok", "content": "hi"}
|
||||
|
||||
result = parse_document(
|
||||
"art-1", {"conversation_id": "c1"}, "user1", {"output": "markdown"},
|
||||
)
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, "art-1", {"conversation_id": "c1"}, "user1",
|
||||
{"output": "markdown"},
|
||||
)
|
||||
assert result == {"status": "ok", "content": "hi"}
|
||||
|
||||
@pytest.mark.unit
|
||||
@patch("application.api.user.tasks.parse_document_worker")
|
||||
def test_soft_time_limit_returns_clean_error(self, mock_worker):
|
||||
from celery.exceptions import SoftTimeLimitExceeded
|
||||
|
||||
from application.api.user.tasks import parse_document
|
||||
|
||||
mock_worker.side_effect = SoftTimeLimitExceeded("parse")
|
||||
|
||||
# The task must swallow the soft-limit signal and return the worker's
|
||||
# clean error shape so the parsing-worker slot frees instead of crashing.
|
||||
result = parse_document("art-1", {"conversation_id": "c1"}, "user1")
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert "timed out" in result["error"]
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_time_limits_derived_from_document_parse_timeout(self):
|
||||
from application.api.user.tasks import parse_document
|
||||
from application.core.settings import settings
|
||||
|
||||
assert parse_document.soft_time_limit == settings.DOCUMENT_PARSE_TIMEOUT
|
||||
assert parse_document.time_limit == settings.DOCUMENT_PARSE_TIMEOUT + 30
|
||||
# Hard limit must exceed the soft limit so the handler can unwind first.
|
||||
assert parse_document.time_limit > parse_document.soft_time_limit
|
||||
|
||||
|
||||
class TestDurableTaskRetryPolicy:
|
||||
"""The long-running tasks share a uniform retry policy."""
|
||||
|
||||
|
||||
@@ -759,6 +759,66 @@ class TestHandleToolCalls:
|
||||
assistant_ids = {tc["id"] for tc in assistant_msg["tool_calls"]}
|
||||
assert tool_msg["tool_call_id"] in assistant_ids
|
||||
|
||||
def test_create_tool_message_failure_no_duplicate_assistant(self):
|
||||
"""A tool that succeeds but returns an unserializable result makes
|
||||
create_tool_message raise *after* the success path appended the
|
||||
assistant tool_calls message. The except branch must reuse that
|
||||
message, not append a second one — exactly one assistant(tool_calls)
|
||||
precedes the role:"tool" error, or the next completion 400s."""
|
||||
|
||||
class _SerializingHandler(ConcreteHandler):
|
||||
def create_tool_message(
|
||||
self, tool_call: ToolCall, result: Any
|
||||
) -> Dict:
|
||||
# Mirror real handlers: JSON-encode non-str results, which
|
||||
# raises on a value the encoder can't serialize. The error
|
||||
# path passes a plain string, so only the success path trips.
|
||||
import json as _json
|
||||
|
||||
content = (
|
||||
result
|
||||
if isinstance(result, str)
|
||||
else _json.dumps(result)
|
||||
)
|
||||
return {
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": content,
|
||||
}
|
||||
|
||||
handler = _SerializingHandler()
|
||||
agent = self._make_agent()
|
||||
agent.tool_executor._name_to_tool = {}
|
||||
|
||||
# Tool runs fine but returns a value json.dumps can't encode.
|
||||
def fake_execute(tools_dict, call):
|
||||
yield {"type": "tool_call", "data": {"status": "pending"}}
|
||||
return (object(), call.id)
|
||||
|
||||
agent._execute_tool_action = Mock(side_effect=fake_execute)
|
||||
|
||||
call = ToolCall(id="c1", name="action_1", arguments="{}")
|
||||
gen = handler.handle_tool_calls(agent, [call], {"1": {"name": "t"}}, [])
|
||||
try:
|
||||
while True:
|
||||
next(gen)
|
||||
except StopIteration as e:
|
||||
messages, _pending = e.value
|
||||
|
||||
# Exactly one assistant(tool_calls) for c1, then the tool error — the
|
||||
# buggy path produced [assistant, assistant, tool].
|
||||
assert [m["role"] for m in messages] == ["assistant", "tool"]
|
||||
assistant_msgs = [
|
||||
m for m in messages
|
||||
if m.get("role") == "assistant"
|
||||
and any(tc["id"] == "c1" for tc in m.get("tool_calls", []))
|
||||
]
|
||||
assert len(assistant_msgs) == 1
|
||||
tool_msg = messages[-1]
|
||||
assert tool_msg["role"] == "tool"
|
||||
assert tool_msg["tool_call_id"] == "c1"
|
||||
assert "Error executing tool" in tool_msg["content"]
|
||||
|
||||
def test_partial_batch_failure_answers_every_tool_call(self):
|
||||
"""One tool raises, others succeed: every assistant tool_call must
|
||||
get a matching role:"tool" reply (OpenAI rejects unanswered ids)."""
|
||||
|
||||
@@ -121,6 +121,94 @@ def test_fast_engine_uses_legacy_parser(monkeypatch):
|
||||
assert out["content"] == "legacy text"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# single Docling conversion: the default markdown+tables path must not re-convert
|
||||
# ---------------------------------------------------------------------------
|
||||
@pytest.mark.unit
|
||||
def test_default_markdown_with_tables_converts_docling_once(monkeypatch):
|
||||
"""Default read (markdown, engine=auto, include_tables) converts the doc with Docling once.
|
||||
|
||||
A Docling-backed parser already converts the whole document to produce its text;
|
||||
collecting tables must reuse that single conversion instead of re-running
|
||||
DocumentConverter. Counts conversions across both sites to prove no double pass.
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
|
||||
from application.parser.file.docling_parser import DoclingParser
|
||||
|
||||
counter = {"instances": 0, "converts": 0}
|
||||
|
||||
class _FakeTable:
|
||||
def export_to_dataframe(self):
|
||||
raise RuntimeError("no dataframe") # force the markdown fallback path
|
||||
|
||||
def export_to_markdown(self):
|
||||
return "| h |\n| - |\n| v |"
|
||||
|
||||
class _FakeDoc:
|
||||
tables = [_FakeTable()]
|
||||
pages = {"1": {}}
|
||||
|
||||
def export_to_markdown(self):
|
||||
return "# single-pass content"
|
||||
|
||||
def export_to_dict(self):
|
||||
return {"texts": [{}], "tables": [{}], "pages": {"1": {}}}
|
||||
|
||||
class _FakeResult:
|
||||
document = _FakeDoc()
|
||||
|
||||
class _CountingConverter:
|
||||
def __init__(self, *args, **kwargs):
|
||||
counter["instances"] += 1
|
||||
|
||||
def convert(self, *args, **kwargs):
|
||||
counter["converts"] += 1
|
||||
return _FakeResult()
|
||||
|
||||
fake_docling = types.ModuleType("docling")
|
||||
fake_dc_module = types.ModuleType("docling.document_converter")
|
||||
fake_dc_module.DocumentConverter = _CountingConverter
|
||||
fake_docling.document_converter = fake_dc_module
|
||||
monkeypatch.setitem(sys.modules, "docling", fake_docling)
|
||||
monkeypatch.setitem(sys.modules, "docling.document_converter", fake_dc_module)
|
||||
|
||||
class _FakeDoclingParser(DoclingParser):
|
||||
"""Docling-backed parser exposing its configured converter + export for reuse.
|
||||
|
||||
The collapse path converts ONCE via ``self._converter`` and exports content via
|
||||
``_export_content`` (the configured pipeline/OCR), so both content and tables
|
||||
come from a single conversion.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._parser_config = {"ready": True} # makes parser_config_set True
|
||||
self._converter = _CountingConverter() # single configured conversion
|
||||
|
||||
def _export_content(self, document):
|
||||
return "# single-pass content"
|
||||
|
||||
def parse_file(self, file, errors="ignore"):
|
||||
# Only reached if the collapse path fails; the test asserts it doesn't.
|
||||
return "# parse-file content"
|
||||
|
||||
monkeypatch.setattr(dr, "get_default_file_extractor", lambda ocr_enabled=None: {".pdf": _FakeDoclingParser()})
|
||||
|
||||
# Defaults: output=markdown, engine=auto, include_tables=True -> the double-parse path.
|
||||
out = parse_document_bytes(b"%PDF-1.4", "doc.pdf")
|
||||
|
||||
assert out["output"] == "markdown"
|
||||
# Content comes from the CONFIGURED parser's export (matches the legacy single
|
||||
# parse), not a vanilla converter; the collapse path was taken, not the fallback.
|
||||
assert "single-pass content" in out["content"]
|
||||
assert "parse-file content" not in out["content"]
|
||||
assert out["tables"] == [{"markdown": "| h |\n| - |\n| v |"}]
|
||||
assert counter["converts"] == 1 # was 2 before the fix (content pass + tables pass)
|
||||
assert counter["instances"] == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pages: page-range slice on a form-feed delimited blob
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -422,3 +422,37 @@ class TestCascadeDelete:
|
||||
)
|
||||
assert repo.get_artifact(created["id"]) is None
|
||||
assert repo.list_versions(created["id"]) == []
|
||||
|
||||
|
||||
class TestDeleteArtifact:
|
||||
def test_delete_removes_versions_and_frees_quota(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
art = repo.create_artifact(
|
||||
"u-del", "document", conversation_id=_conversation_id(),
|
||||
storage_path="k/v1.bin", size=100, filename="a.bin",
|
||||
mime_type="application/octet-stream",
|
||||
)
|
||||
repo.append_version(art["id"], storage_path="k/v2.bin", size=200, filename="a.bin")
|
||||
assert repo.count_for_user("u-del") == 1
|
||||
assert repo.total_bytes_for_user("u-del") == 300
|
||||
|
||||
paths = repo.delete_artifact(art["id"])
|
||||
|
||||
assert set(paths) == {"k/v1.bin", "k/v2.bin"} # caller reaps these bytes
|
||||
assert repo.get_artifact(art["id"]) is None
|
||||
assert repo.list_versions(art["id"]) == [] # versions cascade
|
||||
assert repo.count_for_user("u-del") == 0 # quota freed
|
||||
assert repo.total_bytes_for_user("u-del") == 0
|
||||
|
||||
def test_delete_for_conversation_scopes_to_that_conversation(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
conv = _conversation_id()
|
||||
repo.create_artifact("u2", "document", conversation_id=conv, storage_path="k/a", size=10, filename="a")
|
||||
repo.create_artifact("u2", "document", conversation_id=conv, storage_path="k/b", size=20, filename="b")
|
||||
other = repo.create_artifact("u2", "document", conversation_id=_conversation_id(), storage_path="k/c", size=30, filename="c")
|
||||
|
||||
n = repo.delete_for_conversation(conv)
|
||||
|
||||
assert n == 2
|
||||
assert repo.get_artifact(other["id"]) is not None # other conversation untouched
|
||||
assert repo.count_for_user("u2") == 1
|
||||
Reference in new issue
Block a user