mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
feat(projects): search shelf documents by keyword (#7286)
* feat(projects): search shelf documents by keyword Add search_project_documents: two-tier name/summary then bounded content scan, case-insensitive substring matching with original-text offsets. Never converts on the fly; reports stop_reason, content_scanned, skipped reasons, and capped unscanned/truncated ids. * fix(projects): derive unscanned count without paging after budget trip Review feedback: once the content budget trips, stop paging list_active to enumerate unscanned ids — cap ids at the trip page and derive the honest total from one count_active query (tier 1 is exhaustive, so active - tier1 hits - processed = unscanned). Also add a tool-level isolation test pinning that search never crosses projects or users. * fix(projects): scan the shelf over one stable snapshot, not offset pages Review feedback: offset paging re-read the live shelf per page, so a concurrent upload/trash between pages could duplicate a match or silently skip one while still claiming complete coverage. Tier 1 now materializes the whole shelf in a single SELECT (one read snapshot) and tier 2 iterates those same rows; the budget-trip unscanned count is derived exactly from the snapshot with no extra query. Mid-scan upload/trash regression tests pin both failure modes.
This commit is contained in:
Zeren Wang
GitHub
parent
9b89c7569d
commit
74d2565f64
@ -102,6 +102,18 @@ class ProjectsConfig(BaseModel):
|
||||
field_doc="Caller-side wait limit for one summary model call; abandons the wait (not the remote work), and never covers document conversion.",
|
||||
),
|
||||
)
|
||||
search_max_chars_per_document: int = Field(
|
||||
default=1_000_000,
|
||||
ge=1000,
|
||||
le=50_000_000,
|
||||
description="Per-document scan cap for search_project_documents, in decoded CHARACTERS (not bytes — CJK differs 3x). The tool reads it per call, so edits apply without a restart.",
|
||||
)
|
||||
search_max_total_chars: int = Field(
|
||||
default=8_000_000,
|
||||
ge=1000,
|
||||
le=500_000_000,
|
||||
description=("Whole-shelf total content-scan budget for search_project_documents, in decoded characters; must be >= search_max_chars_per_document (a violation resets both keys to defaults). Read per call; no restart needed."),
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
@ -156,3 +168,21 @@ class ProjectsConfig(BaseModel):
|
||||
else:
|
||||
cleaned[name] = coerced
|
||||
return cleaned
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _search_budget_combo(self) -> ProjectsConfig:
|
||||
"""Combined search-budget rule: total must cover the per-document cap.
|
||||
|
||||
A violation resets BOTH keys together — resetting only one could
|
||||
still leave the combination invalid.
|
||||
"""
|
||||
if self.search_max_total_chars < self.search_max_chars_per_document:
|
||||
logger.warning(
|
||||
"projects.search_max_total_chars (%d) < search_max_chars_per_document (%d); resetting both to defaults",
|
||||
self.search_max_total_chars,
|
||||
self.search_max_chars_per_document,
|
||||
)
|
||||
defaults = type(self)()
|
||||
self.search_max_chars_per_document = defaults.search_max_chars_per_document
|
||||
self.search_max_total_chars = defaults.search_max_total_chars
|
||||
return self
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
"""Project shelf tools: bounded live reads through the pinned project identity.
|
||||
|
||||
Two read-only tools (Phase-2 spec §7.3; the shelf is user-curated — there is
|
||||
no agent-initiated shelf write). Both take ``project_id`` from the run's
|
||||
Read-only tools (Phase-2 spec §7.3; the shelf is user-curated — there is
|
||||
no agent-initiated shelf write). All take ``project_id`` from the run's
|
||||
admission-pinned context (``PROJECT_CONTEXT_KEY``) and ``user_id`` from
|
||||
:func:`resolve_runtime_user_id`, then query **live** shelf rows: the pinned
|
||||
snapshot fixes *which* project, never *what* the shelf currently holds. A
|
||||
@ -19,6 +19,7 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Annotated, Any
|
||||
|
||||
from langchain.tools import tool
|
||||
@ -37,6 +38,14 @@ _LIST_DEFAULT_LIMIT = 50
|
||||
_LIST_MAX_LIMIT = 200
|
||||
_READ_DEFAULT_LIMIT = 8000
|
||||
_READ_MAX_LIMIT = 20000
|
||||
_SEARCH_DEFAULT_LIMIT = 10
|
||||
_SEARCH_MAX_LIMIT = 50
|
||||
_SEARCH_QUERY_MAX_CHARS = 200
|
||||
#: Excerpt window: this many characters on each side of the match (spec §7;
|
||||
#: a constant, deliberately not a config knob).
|
||||
_SEARCH_EXCERPT_CONTEXT = 80
|
||||
#: Cap on each id list in the search result payload.
|
||||
_SEARCH_ID_LIST_CAP = 20
|
||||
|
||||
_BINARY_DECLINE_MESSAGE = "This document is binary and cannot be read as text. To let the agent process it, attach it to a thread (attach-to-thread) instead."
|
||||
_CONVERSION_DISABLED_MESSAGE = "This document is a convertible office/PDF file, but automatic document conversion is disabled (uploads.auto_convert_documents). Enable conversion, or attach the file to a thread (attach-to-thread) instead."
|
||||
@ -204,6 +213,214 @@ async def read_project_document(
|
||||
return await _read_project_document_impl(runtime, document_id=document_id, offset=offset, limit=limit)
|
||||
|
||||
|
||||
_QUERY_BLANK_MESSAGE = "Query must be 1-200 non-whitespace characters."
|
||||
|
||||
|
||||
def _resolve_search_budgets() -> tuple[int, int]:
|
||||
"""Worker-thread config read: per-call so edits apply without a restart."""
|
||||
try:
|
||||
from deerflow.config.app_config import get_app_config
|
||||
from deerflow.config.projects_config import ProjectsConfig
|
||||
|
||||
cfg = getattr(get_app_config(), "projects", None)
|
||||
if not isinstance(cfg, ProjectsConfig):
|
||||
cfg = ProjectsConfig()
|
||||
return cfg.search_max_chars_per_document, cfg.search_max_total_chars
|
||||
except Exception:
|
||||
from deerflow.config.projects_config import ProjectsConfig
|
||||
|
||||
fallback = ProjectsConfig()
|
||||
return fallback.search_max_chars_per_document, fallback.search_max_total_chars
|
||||
|
||||
|
||||
def _search_excerpt(text: str, match: re.Match) -> str:
|
||||
"""Slice ±``_SEARCH_EXCERPT_CONTEXT`` around *match* on the ORIGINAL text, then neutralize.
|
||||
|
||||
Order matters: slicing on neutralized text would misalign every offset.
|
||||
"""
|
||||
start = max(0, match.start() - _SEARCH_EXCERPT_CONTEXT)
|
||||
end = min(len(text), match.end() + _SEARCH_EXCERPT_CONTEXT)
|
||||
snippet = text[start:end]
|
||||
if start > 0:
|
||||
snippet = "…" + snippet
|
||||
if end < len(text):
|
||||
snippet += "…"
|
||||
return neutralize_untrusted_tags(snippet)
|
||||
|
||||
|
||||
def _capped_id_list(ids: list[str]) -> tuple[list[str], int]:
|
||||
"""The first ``_SEARCH_ID_LIST_CAP`` ids plus the full count (spec §6)."""
|
||||
return ids[:_SEARCH_ID_LIST_CAP], len(ids)
|
||||
|
||||
|
||||
async def _search_project_documents_impl(runtime: Runtime, *, query: str, limit: int) -> str:
|
||||
resolved = _resolve_pin_and_repo(runtime)
|
||||
if isinstance(resolved, str):
|
||||
return resolved
|
||||
project_id, user_id, repo = resolved
|
||||
|
||||
query = query.strip()
|
||||
if not query or len(query) > _SEARCH_QUERY_MAX_CHARS:
|
||||
return _error(_QUERY_BLANK_MESSAGE)
|
||||
limit = _clamp(limit, default=_SEARCH_DEFAULT_LIMIT, lower=1, upper=_SEARCH_MAX_LIMIT)
|
||||
per_doc_cap, total_budget = await run_file_io(_resolve_search_budgets)
|
||||
# re.IGNORECASE, NOT casefold: indices must stay original-text character
|
||||
# positions — casefold changes length (ß→ss, İ→i̇) and would corrupt both
|
||||
# the returned offset and the excerpt slice (spec §5.3).
|
||||
pattern = re.compile(re.escape(query), re.IGNORECASE)
|
||||
|
||||
# Tier 1: name/summary over the whole shelf, materialized in ONE query.
|
||||
# The rows are kept and reused for tier 2 — a stable candidate set.
|
||||
# Offset paging would re-read the live shelf per page, so a concurrent
|
||||
# upload/trash between pages could duplicate a row (an unrelated newest
|
||||
# upload shifts every offset) or silently skip one (a trash shrinks
|
||||
# earlier pages); a single SELECT reads one snapshot. A row trashed
|
||||
# between the snapshot and its content read follows the normal serving
|
||||
# rules ("missing" etc.) — shelf membership is fixed as of this call.
|
||||
shelf_rows = await repo.list_active(project_id, limit=1 << 62, offset=0, user_id=user_id) # effectively unbounded: one read snapshot
|
||||
name_hits: list[dict] = []
|
||||
summary_hits: list[dict] = []
|
||||
tier1_ids: set[str] = set()
|
||||
for row in shelf_rows:
|
||||
name = str(row.get("name") or "")
|
||||
summary = row.get("summary") or None
|
||||
name_match = pattern.search(name)
|
||||
if name_match is not None:
|
||||
name_hits.append({"id": row["id"], "name": neutralize_untrusted_tags(name), "matched": "name", "excerpt": _search_excerpt(name, name_match), "offset": None})
|
||||
tier1_ids.add(row["id"])
|
||||
elif summary:
|
||||
summary_match = pattern.search(str(summary))
|
||||
if summary_match is not None:
|
||||
summary_hits.append({"id": row["id"], "name": neutralize_untrusted_tags(name), "matched": "summary", "excerpt": _search_excerpt(str(summary), summary_match), "offset": None})
|
||||
tier1_ids.add(row["id"])
|
||||
|
||||
results = (name_hits + summary_hits)[:limit]
|
||||
if len(results) >= limit:
|
||||
# Tier 1 alone filled the page: content was never read, and the
|
||||
# ranking is exact — the answer must NOT claim coverage (spec §6).
|
||||
return json.dumps(
|
||||
{
|
||||
"results": results,
|
||||
"stop_reason": "limit",
|
||||
"content_scanned": False,
|
||||
"skipped": {"non_text": 0, "unconverted": 0, "missing": 0, "error": 0},
|
||||
"truncated_document_ids": [],
|
||||
"truncated_count": 0,
|
||||
"unscanned_document_ids": [],
|
||||
"unscanned_count": 0,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
|
||||
# Tier 2: content scan over the SAME snapshot, in shelf order, skipping
|
||||
# tier-1 hits, bounded by the whole-shelf budget. auto_convert=False:
|
||||
# search NEVER converts (a first search must not synchronously convert
|
||||
# every PDF on the shelf).
|
||||
remaining = limit - len(results)
|
||||
budget_left = total_budget
|
||||
stop_reason = "complete"
|
||||
skipped = {"non_text": 0, "unconverted": 0, "missing": 0, "error": 0}
|
||||
truncated_ids: list[str] = []
|
||||
unscanned_ids: list[str] = []
|
||||
processed = 0 # non-tier-1 rows examined before any budget trip
|
||||
paths = get_paths()
|
||||
for row in shelf_rows:
|
||||
if remaining == 0:
|
||||
stop_reason = "limit"
|
||||
break
|
||||
if row["id"] in tier1_ids:
|
||||
# Tier 1 was fully scanned over this snapshot: these are already
|
||||
# in results and were searched — never "unscanned".
|
||||
continue
|
||||
if budget_left <= 0:
|
||||
# Budget tripped. Unscanned ids are enumerated in memory, up to
|
||||
# the payload cap; the honest total is derived after the loop
|
||||
# from the snapshot — no further DB reads for rows already ruled
|
||||
# out.
|
||||
stop_reason = "budget"
|
||||
if len(unscanned_ids) < _SEARCH_ID_LIST_CAP:
|
||||
unscanned_ids.append(row["id"])
|
||||
continue
|
||||
processed += 1
|
||||
try:
|
||||
serving_path, reason = await read_text_serving_path(repo, paths, user_id=user_id, row=row, auto_convert=False)
|
||||
except Exception:
|
||||
skipped["error"] += 1
|
||||
logger.warning("Shelf search: failed to resolve serving path for %s", row["id"], exc_info=True)
|
||||
continue
|
||||
if serving_path is None:
|
||||
skipped[{"content_missing": "missing", "conversion_disabled": "unconverted"}.get(reason, "non_text")] += 1
|
||||
continue
|
||||
cap = min(per_doc_cap, budget_left)
|
||||
try:
|
||||
text = await read_document_text_window(serving_path, offset=0, limit=cap + 1)
|
||||
except Exception:
|
||||
skipped["error"] += 1
|
||||
logger.warning("Shelf search: failed to read %s", row["id"], exc_info=True)
|
||||
continue
|
||||
budget_left -= len(text)
|
||||
if len(text) > cap:
|
||||
text = text[:cap]
|
||||
truncated_ids.append(row["id"])
|
||||
match = pattern.search(text)
|
||||
if match is not None:
|
||||
results.append({"id": row["id"], "name": neutralize_untrusted_tags(str(row.get("name") or "")), "matched": "content", "excerpt": _search_excerpt(text, match), "offset": match.start()})
|
||||
remaining -= 1
|
||||
|
||||
truncated_capped, truncated_total = _capped_id_list(truncated_ids)
|
||||
if stop_reason == "complete" and remaining == 0:
|
||||
# The limit filled exactly on the snapshot's final row — a limit stop,
|
||||
# not a complete scan.
|
||||
stop_reason = "limit"
|
||||
unscanned_total = 0
|
||||
if stop_reason == "budget":
|
||||
# Exact — the snapshot IS the candidate set: unscanned = non-tier-1
|
||||
# rows minus those processed before the trip.
|
||||
unscanned_total = len(shelf_rows) - len(tier1_ids) - processed
|
||||
return json.dumps(
|
||||
{
|
||||
"results": results,
|
||||
"stop_reason": stop_reason,
|
||||
"content_scanned": True,
|
||||
"skipped": skipped,
|
||||
"truncated_document_ids": truncated_capped,
|
||||
"truncated_count": truncated_total,
|
||||
"unscanned_document_ids": unscanned_ids,
|
||||
"unscanned_count": unscanned_total,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
|
||||
|
||||
@tool
|
||||
async def search_project_documents(
|
||||
runtime: Runtime,
|
||||
query: Annotated[str, "Case-insensitive keyword to find on the shelf (plain substring, not regex; 1-200 chars)."],
|
||||
limit: Annotated[int, "Maximum results to return (default 10, max 50)."] = _SEARCH_DEFAULT_LIMIT,
|
||||
) -> str:
|
||||
"""Search the current project's shelf by keyword across names, summaries, and content.
|
||||
|
||||
Two tiers: every document's name and summary are matched first (cheap),
|
||||
and content is scanned only if that produced too few hits — so ranking is
|
||||
always name > summary > content, with the shelf's own order inside a
|
||||
tier. Each document appears at most once, credited to its best field.
|
||||
|
||||
Returns JSON: {"results": [{"id", "name", "matched", "excerpt",
|
||||
"offset"}], "stop_reason", "content_scanned", "skipped",
|
||||
"truncated_document_ids", "truncated_count", "unscanned_document_ids",
|
||||
"unscanned_count"}. "offset" is set ONLY for content hits: a character
|
||||
position to pass straight to read_project_document(offset=...).
|
||||
stop_reason is "complete" (everything searched), "limit" (enough hits;
|
||||
unsearched remainder is irrelevant), or "budget" (the whole-shelf scan
|
||||
budget ran out — unscanned_document_ids lists what was not searched).
|
||||
content_scanned=false means content was never read (tier 1 was enough);
|
||||
never treat that as full coverage. skipped counts non_text binaries,
|
||||
unconverted office/PDF (never converted on the fly), and missing
|
||||
content. Live data: reflects the shelf as of this call.
|
||||
"""
|
||||
return await _search_project_documents_impl(runtime, query=query, limit=limit)
|
||||
|
||||
|
||||
def get_project_document_tools() -> list:
|
||||
"""The two shelf tools the lead agent gains in project runs (§7.3)."""
|
||||
return [list_project_documents, read_project_document]
|
||||
"""The shelf tools the lead agent gains in project runs (§7.3)."""
|
||||
return [list_project_documents, read_project_document, search_project_documents]
|
||||
|
||||
@ -206,3 +206,27 @@ async def test_conversion_publish_path_dispatches_off_the_loop(tmp_path, monkeyp
|
||||
assert "_cached_char_count" in calls
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
|
||||
async def test_tool_search_dispatches_off_the_loop(tmp_path, monkeypatch) -> None:
|
||||
"""The content tier of ``search_project_documents`` validates the original
|
||||
(``_content_intact``), resolves serving paths, samples text
|
||||
(``is_text_file_by_content``) and reads the scan window
|
||||
(``_read_text_window``) — all through the offload."""
|
||||
from deerflow.projects.tools import _search_project_documents_impl
|
||||
|
||||
monkeypatch.setattr("deerflow.projects.tools._resolve_search_budgets", lambda: (100_000, 1_000_000))
|
||||
calls = _spy_offload(monkeypatch)
|
||||
env = await _make_env(tmp_path, monkeypatch)
|
||||
try:
|
||||
staged = await stage_document_bytes(env.paths, user_id=_USER, project_id=env.project_id, chunks=[b"needle in the haystack"], max_bytes=1 << 20)
|
||||
await add_staged_document(env.docs, env.paths, user_id=_USER, project_id=env.project_id, name="hay.txt", staged=staged)
|
||||
calls.clear()
|
||||
payload = json.loads(await _search_project_documents_impl(_runtime(env.project_id), query="needle", limit=10))
|
||||
assert payload["results"] and payload["results"][0]["matched"] == "content"
|
||||
assert "_content_intact" in calls
|
||||
assert "resolve_document_paths" in calls
|
||||
assert "is_text_file_by_content" in calls
|
||||
assert "_read_text_window" in calls
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
361
backend/tests/test_project_document_search.py
Normal file
361
backend/tests/test_project_document_search.py
Normal file
@ -0,0 +1,361 @@
|
||||
"""Tests for the shelf search tool (search_project_documents).
|
||||
|
||||
Covers the two-tier ordering, Unicode-offset integrity (re.IGNORECASE, never
|
||||
casefold), the never-convert rule, budget/coverage honesty, skip categories,
|
||||
excerpt neutralization, and query validation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.projects import tools as tools_mod
|
||||
from deerflow.projects.documents import add_staged_document, original_file_path, stage_document_bytes
|
||||
from deerflow.projects.tools import _search_project_documents_impl
|
||||
from deerflow.runtime.context_keys import PROJECT_CONTEXT_KEY
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
_USER = "u1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def env(tmp_path, monkeypatch):
|
||||
"""Real SQLite repos + a real per-user projects layout on disk."""
|
||||
import deerflow.config.paths as paths_mod
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
from deerflow.persistence.projects import ProjectDocumentRepository, ProjectRepository
|
||||
|
||||
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
|
||||
monkeypatch.setattr(paths_mod, "_paths", None)
|
||||
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'test.db'}", sqlite_dir=str(tmp_path))
|
||||
sf = get_session_factory()
|
||||
projects = ProjectRepository(sf)
|
||||
docs = ProjectDocumentRepository(sf)
|
||||
project = await projects.create(name="P", user_id=_USER)
|
||||
yield SimpleNamespace(paths=paths_mod.get_paths(), sf=sf, projects=projects, docs=docs, project=project)
|
||||
await close_engine()
|
||||
|
||||
|
||||
def _runtime(*, project_id: str | None, user_id: str = _USER) -> SimpleNamespace:
|
||||
context: dict = {"user_id": user_id}
|
||||
if project_id is not None:
|
||||
context[PROJECT_CONTEXT_KEY] = {"project_id": project_id, "name": "P", "instructions": ""}
|
||||
return SimpleNamespace(context=context)
|
||||
|
||||
|
||||
async def _shelve(env, *, name: str, data: bytes, project_id: str | None = None, user_id: str = _USER) -> dict:
|
||||
staged = await stage_document_bytes(env.paths, user_id=user_id, project_id=project_id or env.project["id"], chunks=[data], max_bytes=1 << 20)
|
||||
result = await add_staged_document(env.docs, env.paths, user_id=user_id, project_id=project_id or env.project["id"], name=name, staged=staged)
|
||||
assert result is not None
|
||||
row, created = result
|
||||
assert created
|
||||
return row
|
||||
|
||||
|
||||
def _budgets(monkeypatch, per_doc: int, total: int) -> None:
|
||||
monkeypatch.setattr(tools_mod, "_resolve_search_budgets", lambda: (per_doc, total))
|
||||
|
||||
|
||||
async def _search(env, query: str, limit: int = 10) -> dict:
|
||||
return json.loads(await _search_project_documents_impl(_runtime(project_id=env.project["id"]), query=query, limit=limit))
|
||||
|
||||
|
||||
class TestTierOrdering:
|
||||
async def test_name_hit_ranks_before_content_hit(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="other.txt", data=b"the keyword hides here")
|
||||
await _shelve(env, name="keyword-report.txt", data=b"nothing relevant")
|
||||
result = await _search(env, "keyword")
|
||||
assert [r["matched"] for r in result["results"]] == ["name", "content"]
|
||||
assert result["stop_reason"] == "complete"
|
||||
assert result["content_scanned"] is True
|
||||
|
||||
async def test_summary_hit_ranks_before_content_hit(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
row = await _shelve(env, name="plain.txt", data=b"nothing")
|
||||
assert await env.docs.set_summary(row["id"], "keyword 摘要", user_id=_USER)
|
||||
await _shelve(env, name="content.txt", data=b"keyword in body")
|
||||
result = await _search(env, "keyword")
|
||||
assert [r["matched"] for r in result["results"]] == ["summary", "content"]
|
||||
|
||||
async def test_tier1_filling_limit_never_reads_content(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
reads = 0
|
||||
|
||||
async def counting_window(path, *, offset, limit):
|
||||
nonlocal reads
|
||||
reads += 1
|
||||
return ""
|
||||
|
||||
monkeypatch.setattr(tools_mod, "read_document_text_window", counting_window)
|
||||
for i in range(3):
|
||||
await _shelve(env, name=f"hit-{i}.txt", data=f"irrelevant {i}".encode())
|
||||
result = await _search(env, "hit", limit=2)
|
||||
assert result["stop_reason"] == "limit"
|
||||
assert result["content_scanned"] is False
|
||||
assert len(result["results"]) == 2
|
||||
assert reads == 0
|
||||
|
||||
async def test_same_document_matches_once(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="dup.txt", data=b"hit and hit and hit")
|
||||
result = await _search(env, "hit")
|
||||
assert len(result["results"]) == 1
|
||||
|
||||
|
||||
class TestUnicodeOffsets:
|
||||
async def test_ignorecase_offsets_point_at_original_text(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
body = "prefix İstanbul suffix"
|
||||
await _shelve(env, name="city.txt", data=body.encode())
|
||||
result = await _search(env, "istanbul")
|
||||
hit = result["results"][0]
|
||||
assert hit["matched"] == "content"
|
||||
# re.IGNORECASE keeps original indices — the offset must index the
|
||||
# original text exactly (casefold would corrupt this).
|
||||
assert hit["offset"] == body.index("İstanbul")
|
||||
|
||||
async def test_multichar_fold_equivalents_do_not_match(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
body = "die Straße ist lang"
|
||||
await _shelve(env, name="street.txt", data=body.encode())
|
||||
# ß↔ss is a multi-char fold equivalence — deliberately NOT matched.
|
||||
assert (await _search(env, "strasse"))["results"] == []
|
||||
result = await _search(env, "straße")
|
||||
assert result["results"][0]["offset"] == body.index("Straße")
|
||||
|
||||
async def test_offset_feeds_read_project_document(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
body = "aaaa bbbb needle zzzz"
|
||||
row = await _shelve(env, name="doc.txt", data=body.encode())
|
||||
result = await _search(env, "needle")
|
||||
offset = result["results"][0]["offset"]
|
||||
from deerflow.projects.tools import _read_project_document_impl
|
||||
|
||||
page = json.loads(await _read_project_document_impl(_runtime(project_id=env.project["id"]), document_id=row["id"], offset=offset, limit=20))
|
||||
assert page["content"].startswith("needle")
|
||||
|
||||
|
||||
class TestNeverConvert:
|
||||
async def test_unconverted_convertible_is_skipped_not_converted(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
calls = 0
|
||||
|
||||
async def counting_convert(file_path, output_path=None):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return ""
|
||||
|
||||
from deerflow.projects import documents as documents_mod
|
||||
|
||||
monkeypatch.setattr(documents_mod, "convert_file_to_markdown", counting_convert)
|
||||
await _shelve(env, name="slides.pdf", data=b"%PDF-1.4 fake")
|
||||
result = await _search(env, "fake")
|
||||
assert calls == 0
|
||||
assert result["skipped"]["unconverted"] == 1
|
||||
|
||||
async def test_binary_and_missing_are_counted_not_fatal(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="blob.bin", data=b"\x00\x01\x02hit\x00")
|
||||
lost = await _shelve(env, name="lost.txt", data=b"hit was here")
|
||||
original_file_path(env.paths, user_id=_USER, row=lost).unlink()
|
||||
result = await _search(env, "hit")
|
||||
assert result["skipped"]["non_text"] == 1
|
||||
assert result["skipped"]["missing"] == 1
|
||||
assert result["skipped"]["error"] == 0
|
||||
|
||||
|
||||
class TestBudgetsAndCoverage:
|
||||
async def test_budget_exhaustion_reports_unscanned(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 100, 150)
|
||||
for i in range(3):
|
||||
await _shelve(env, name=f"big-{i}.txt", data=(f"x{i}" * 250 + " needle").encode())
|
||||
result = await _search(env, "needle")
|
||||
assert result["stop_reason"] == "budget"
|
||||
assert result["truncated_count"] == 2
|
||||
assert result["unscanned_count"] == 1
|
||||
|
||||
async def test_per_doc_truncation_marks_only_that_document(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10, 10_000)
|
||||
long_row = await _shelve(env, name="long.txt", data=b"needle " + b"y" * 500)
|
||||
await _shelve(env, name="short.txt", data=b"ok")
|
||||
result = await _search(env, "needle")
|
||||
assert result["truncated_document_ids"] == [long_row["id"]]
|
||||
assert result["truncated_count"] == 1
|
||||
assert result["stop_reason"] == "complete"
|
||||
|
||||
async def test_tier1_hit_after_budget_exhaustion_is_not_unscanned(self, env, monkeypatch):
|
||||
"""A name hit sitting after the budget trip point was already searched
|
||||
(tier 1 completes before tier 2 starts) — it must not be reported as
|
||||
unscanned even though content scanning never reached it."""
|
||||
_budgets(monkeypatch, 100, 100)
|
||||
# Shelf order is updated_at DESC: insert order is reversed below, so
|
||||
# the content doc scans first, big2 trips the budget, and the name hit
|
||||
# is encountered only afterwards.
|
||||
name_hit = await _shelve(env, name="needle-report.txt", data=b"irrelevant")
|
||||
big2 = await _shelve(env, name="big2.txt", data=b"y" * 250)
|
||||
big1 = await _shelve(env, name="big1.txt", data=b"needle " + b"x" * 250)
|
||||
result = await _search(env, "needle")
|
||||
assert result["stop_reason"] == "budget"
|
||||
assert {r["id"] for r in result["results"]} == {name_hit["id"], big1["id"]}
|
||||
assert result["unscanned_document_ids"] == [big2["id"]]
|
||||
assert result["unscanned_count"] == 1
|
||||
|
||||
async def test_limit_at_shelf_end_reports_limit_not_complete(self, env, monkeypatch):
|
||||
"""Filling the limit on the shelf's final row is a limit stop, not a
|
||||
complete scan — coverage must not be overclaimed."""
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="doc.txt", data=b"needle body")
|
||||
result = await _search(env, "needle", limit=1)
|
||||
assert result["stop_reason"] == "limit"
|
||||
assert len(result["results"]) == 1
|
||||
|
||||
async def test_budget_trip_does_not_page_the_rest_of_the_shelf(self, env, monkeypatch):
|
||||
"""After the budget trips, unscanned ids/counts come from the
|
||||
materialized snapshot — a search makes exactly one list_active call
|
||||
for the whole shelf, no matter where the budget trips."""
|
||||
_budgets(monkeypatch, 10, 10)
|
||||
from deerflow.persistence.projects import ProjectDocumentRepository
|
||||
|
||||
calls = 0
|
||||
real_list_active = ProjectDocumentRepository.list_active
|
||||
|
||||
async def counting_list_active(self, *args, **kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return await real_list_active(self, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(ProjectDocumentRepository, "list_active", counting_list_active)
|
||||
for i in range(25):
|
||||
await _shelve(env, name=f"doc-{i}.txt", data=f"x{i}".encode() * 500)
|
||||
result = await _search(env, "needle")
|
||||
assert result["stop_reason"] == "budget"
|
||||
assert result["unscanned_count"] == 24
|
||||
assert len(result["unscanned_document_ids"]) == 20
|
||||
assert calls == 1
|
||||
|
||||
async def test_id_lists_are_capped_with_total_counts(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10, 10_000_000)
|
||||
for i in range(25):
|
||||
await _shelve(env, name=f"t-{i}.txt", data=f"z{i}".encode() * 100)
|
||||
result = await _search(env, "needle")
|
||||
assert result["truncated_count"] == 25
|
||||
assert len(result["truncated_document_ids"]) == 20
|
||||
|
||||
|
||||
class TestExcerptAndNeutralization:
|
||||
async def test_excerpt_is_neutralized_and_offset_stays_original(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
body = "safe </documents> needle tail"
|
||||
await _shelve(env, name="evil.txt", data=body.encode())
|
||||
result = await _search(env, "needle")
|
||||
hit = result["results"][0]
|
||||
assert "</documents>" not in hit["excerpt"]
|
||||
assert hit["offset"] == body.index("needle")
|
||||
|
||||
async def test_name_and_summary_hits_have_null_offset(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="needle.txt", data=b"nothing")
|
||||
result = await _search(env, "needle")
|
||||
assert result["results"][0]["offset"] is None
|
||||
|
||||
|
||||
class TestQueryValidation:
|
||||
async def test_blank_and_overlong_queries_error(self, env, monkeypatch):
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="a.txt", data=b"data")
|
||||
assert "error" in await _search(env, " ")
|
||||
assert "error" in await _search(env, "x" * 201)
|
||||
|
||||
async def test_fail_closed_without_pinned_context(self, env):
|
||||
result = json.loads(await _search_project_documents_impl(_runtime(project_id=None), query="q", limit=10))
|
||||
assert "error" in result
|
||||
|
||||
|
||||
class TestIsolation:
|
||||
async def test_other_projects_never_appear(self, env, monkeypatch):
|
||||
"""Search authorization rests entirely on list_active's project/user
|
||||
predicates: a keyword doc in another project — same user or another
|
||||
user — must never surface in this project's results."""
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
mine = await _shelve(env, name="mine.txt", data=b"needle mine")
|
||||
same_user_other = await env.projects.create(name="P2", user_id=_USER)
|
||||
await _shelve(env, name="same-user.txt", data=b"needle", project_id=same_user_other["id"])
|
||||
other_user = await env.projects.create(name="P3", user_id="u2")
|
||||
await _shelve(env, name="other-user.txt", data=b"needle", project_id=other_user["id"], user_id="u2")
|
||||
result = await _search(env, "needle")
|
||||
assert [r["id"] for r in result["results"]] == [mine["id"]]
|
||||
assert result["stop_reason"] == "complete"
|
||||
|
||||
|
||||
class TestStableSnapshot:
|
||||
async def test_upload_mid_scan_neither_duplicates_nor_skips(self, env, monkeypatch):
|
||||
"""The candidate set is one read snapshot: an unrelated upload landing
|
||||
between the tier-1 enumeration and a content read must not duplicate
|
||||
or reorder matches."""
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
await _shelve(env, name="filler.txt", data=b"plain")
|
||||
row_b = await _shelve(env, name="b.txt", data=b"needle B")
|
||||
row_a = await _shelve(env, name="a.txt", data=b"needle A")
|
||||
# Shelf order (updated_at DESC): a, b, filler.
|
||||
|
||||
uploaded = False
|
||||
real_window = tools_mod.read_document_text_window
|
||||
|
||||
async def mutating_window(path, *, offset, limit):
|
||||
nonlocal uploaded
|
||||
if not uploaded:
|
||||
uploaded = True
|
||||
await _shelve(env, name="z-new.txt", data=b"unrelated newest upload")
|
||||
return await real_window(path, offset=offset, limit=limit)
|
||||
|
||||
monkeypatch.setattr(tools_mod, "read_document_text_window", mutating_window)
|
||||
result = await _search(env, "needle", limit=2)
|
||||
assert [r["id"] for r in result["results"]] == [row_a["id"], row_b["id"]]
|
||||
assert result["stop_reason"] == "limit"
|
||||
|
||||
async def test_trash_mid_scan_neither_skips_nor_hides_coverage(self, env, monkeypatch):
|
||||
"""Trashing an already-scanned nonmatching document mid-scan must not
|
||||
make the scan skip a later match while still claiming complete
|
||||
coverage."""
|
||||
_budgets(monkeypatch, 10_000, 100_000)
|
||||
row_a = await _shelve(env, name="a.txt", data=b"needle A")
|
||||
filler = await _shelve(env, name="filler.txt", data=b"plain")
|
||||
row_b = await _shelve(env, name="b.txt", data=b"needle B")
|
||||
# Shelf order (updated_at DESC): b, filler, a.
|
||||
|
||||
trashed = False
|
||||
real_window = tools_mod.read_document_text_window
|
||||
|
||||
async def mutating_window(path, *, offset, limit):
|
||||
nonlocal trashed
|
||||
if not trashed:
|
||||
trashed = True
|
||||
assert await env.docs.trash(filler["id"], user_id=_USER)
|
||||
return await real_window(path, offset=offset, limit=limit)
|
||||
|
||||
monkeypatch.setattr(tools_mod, "read_document_text_window", mutating_window)
|
||||
result = await _search(env, "needle", limit=2)
|
||||
assert [r["id"] for r in result["results"]] == [row_b["id"], row_a["id"]]
|
||||
assert result["stop_reason"] == "limit"
|
||||
assert result["unscanned_count"] == 0
|
||||
|
||||
|
||||
class TestConfig:
|
||||
def test_search_budget_defaults_and_bounds(self):
|
||||
from deerflow.config.projects_config import ProjectsConfig
|
||||
|
||||
cfg = ProjectsConfig()
|
||||
assert cfg.search_max_chars_per_document == 1_000_000
|
||||
assert cfg.search_max_total_chars == 8_000_000
|
||||
|
||||
def test_combo_violation_resets_both_keys(self):
|
||||
from deerflow.config.projects_config import ProjectsConfig
|
||||
|
||||
cfg = ProjectsConfig(search_max_chars_per_document=5000, search_max_total_chars=1000)
|
||||
assert cfg.search_max_chars_per_document == 1_000_000
|
||||
assert cfg.search_max_total_chars == 8_000_000
|
||||
@ -20,7 +20,7 @@
|
||||
# ============================================================================
|
||||
# Bump this number when the config schema changes.
|
||||
# Run `make config-upgrade` to merge new fields into your local config.yaml.
|
||||
config_version: 61
|
||||
config_version: 62
|
||||
|
||||
# Optional operator-owned literal prompt extensions (see CONFIGURATION.md):
|
||||
# lead_prompt_overlay:
|
||||
@ -1647,6 +1647,12 @@ projects:
|
||||
# Caller-side wait limit in seconds for one summary model call (range:
|
||||
# 5-300); abandons the wait, not the remote work.
|
||||
summary_timeout_seconds: 60
|
||||
# Shelf search budgets (search_project_documents), in decoded CHARACTERS
|
||||
# not bytes (CJK differs 3x). Per-document cap and whole-shelf total; the
|
||||
# total must be >= the per-document cap, otherwise both reset to defaults.
|
||||
# Read per call — edits apply without a restart.
|
||||
search_max_chars_per_document: 1000000
|
||||
search_max_total_chars: 8000000
|
||||
|
||||
# ============================================================================
|
||||
# Provider Safety Termination Configuration
|
||||
|
||||
@ -135,7 +135,7 @@ they resolve from the `secrets` map):
|
||||
|
||||
```yaml
|
||||
config: |
|
||||
config_version: 61
|
||||
config_version: 62
|
||||
models:
|
||||
- name: gpt-4
|
||||
use: langchain_openai:ChatOpenAI
|
||||
|
||||
@ -367,7 +367,7 @@ ingress:
|
||||
# -- DeerFlow config.yaml content. Secrets MUST stay as $VAR references — never
|
||||
# inline literal secret values here. The default enables provisioner sandbox.
|
||||
config: |
|
||||
config_version: 61
|
||||
config_version: 62
|
||||
log_level: info
|
||||
recursion_limit: 100
|
||||
max_recursion_limit: 1000
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user