mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-06 02:15:54 +00:00
Alex
795e39a6bc
Source access control --------------------- `active_docs` is client-supplied and reached the retriever unchecked, and the retriever queries `WHERE source_id = <id>` with no owner predicate — so any caller could pass any source id to /stream or /api/answer and have another tenant's documents quoted back, while /api/sources/<id>/search correctly refused the same id. Gate it through `can_access`, the helper the guarded endpoints already use, and filter `self.source` down to the authorized set. Fails closed: no principal, or a check that errors, drops the source. Three sibling paths had the same gap: - workflow agent nodes: `AgentNodeConfig.sources` is written verbatim from client JSON at save time and nothing validated it, so a node could name any tenant's source. Gate against the workflow owner, so shared workflows keep reading their owner's sources like shared agents do. - /api/share: `_resolve_source_pg_id` resolved any id with no ownership predicate and baked it into the agent the share creates; /api/search then searched it. Authorize before attaching. - search_service: re-resolve the ids stored on an agent row instead of trusting them, so a row written by any future path with the same gap cannot be read back. Team grantees previously lost their source's retrieval config: the post-check read was still owner-scoped, so it missed and fell back to defaults (an `agentic_tool` source was bulk-prefetched for every grantee). Read unscoped after `can_access` passes. Retrieval --------- `PGVectorStore._ensure_table_exists` created an IVFFlat index on the empty table it had just created. IVFFlat computes centroids at build time, so those centroids were random, and combined with the `source_id` post-filter a source with hundreds of embedded chunks returned zero rows — retrieval reported no documents, the model answered from memory, and nothing was logged. Stop creating the index (exact search is correct and fast well past the sizes most deployments reach); raise `ivfflat.probes` to sqrt(lists) where an index still exists; and re-run a short indexed search exactly, since post-filtering means no index setting can guarantee a full result. `graphrag` had the same empty-table index with no fallback at all. Also: bound `chunks` to 0-500 on both the request and agent paths (0 still means "skip retrieval"), let a source's configured `retrieval.chunks` outrank the request body, and cap ClassicRAG's per-source floor at max(top_k, n_sources) so attaching sources cannot inflate the result set. Silent failures --------------- An empty retrieval was invisible to both the model and the client: the `source` event was suppressed when the list was empty, so "searched and found nothing" looked identical to "no source attached", and the prompt said nothing at all. Emit the event always, and tell the model when a search ran and returned nothing. A file that parses to nothing now fails ingest with a message naming the cause instead of storing an embedding of the empty string. `score_threshold` returns warnings when the active store or retriever cannot honour it. Prompt structure ---------------- Retrieved documents move from the system prompt into the user turn, with the injection guard restated next to them: they change every turn (defeating prefix caching), they are third-party text that should not carry system authority, and routing them through the query budget makes them truncatable rather than silently crowding it out. Documents are shed lowest-ranked-first before the question is touched. The six chat presets (3 tones x 2 retrieval modes) differed only in their Answering section; they are now composed from single-source fragments at load time, not through Jinja inheritance, which would have opened a file-read surface in the template sandbox and broken the tool-prefetch parser. Per-tool guidance moves out of the prompt into tool schemas, so it travels with the tool and cannot render when the tool is absent. A plain-text custom prompt is staged as a persona value inside the skeleton instead of replacing it wholesale — it used to silently lose the injection guard, platform block, memory and attachments, and its braces are now inert. Other fixes ----------- - agents/base: an oversized system prompt drove the query budget negative and dispatched a full-price request with an empty question; raise instead. - llm/anthropic: migrate off the retired Text Completions API. It flattened history to first+last message and ignored tools entirely. Adds the missing Anthropic handler, without which every tool call was silently dropped. - sources/upload: `sitemap` had no branch, so every sitemap ingest died on a TypeError; `validate_url` now rejects a falsy URL cleanly. - workflow nodes: retrieved documents never reached the node agent, so a classic node with a source and an ordinary prompt answered "I have no documents" while the run reported completed. - parser/bulk: copy the metadata dict, or every chunk reports the last chunk's token_count. - crawler_loader: carry the page title, or citations render the whole chunk body as the label.
571 lines
18 KiB
Python
571 lines
18 KiB
Python
import uuid
|
|
from contextlib import contextmanager
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
|
|
|
|
@pytest.fixture
|
|
def app():
|
|
app = Flask(__name__)
|
|
return app
|
|
|
|
|
|
@contextmanager
|
|
def _patch_db(conn):
|
|
"""Patch both db_session and db_readonly to yield the given conn."""
|
|
@contextmanager
|
|
def _yield_conn():
|
|
yield conn
|
|
|
|
with patch(
|
|
"application.api.user.prompts.routes.db_session", _yield_conn
|
|
), patch(
|
|
"application.api.user.prompts.routes.db_readonly", _yield_conn
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCreatePrompt:
|
|
pass
|
|
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from application.api.user.prompts.routes import CreatePrompt
|
|
|
|
with app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "P", "content": "C"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = None
|
|
response = CreatePrompt().post()
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_returns_400_missing_fields(self, app):
|
|
from application.api.user.prompts.routes import CreatePrompt
|
|
|
|
with app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "P"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = CreatePrompt().post()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestGetPrompts:
|
|
pass
|
|
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from application.api.user.prompts.routes import GetPrompts
|
|
|
|
with app.test_request_context("/api/get_prompts"):
|
|
from flask import request
|
|
|
|
request.decoded_token = None
|
|
response = GetPrompts().get()
|
|
|
|
assert response.status_code == 401
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestGetSinglePrompt:
|
|
pass
|
|
|
|
def test_returns_default_prompt(self, app):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
from application.prompts.composer import compose_preset
|
|
|
|
with app.test_request_context("/api/get_single_prompt?id=default"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["content"] == compose_preset("default")
|
|
|
|
def test_returns_creative_prompt(self, app):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
from application.prompts.composer import compose_preset
|
|
|
|
with app.test_request_context("/api/get_single_prompt?id=creative"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["content"] == compose_preset("creative")
|
|
|
|
def test_returns_strict_prompt(self, app):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
from application.prompts.composer import compose_preset
|
|
|
|
with app.test_request_context("/api/get_single_prompt?id=strict"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["content"] == compose_preset("strict")
|
|
|
|
|
|
def test_returns_400_missing_id(self, app):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
with app.test_request_context("/api/get_single_prompt"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestDeletePrompt:
|
|
pass
|
|
|
|
def test_returns_400_missing_id(self, app):
|
|
from application.api.user.prompts.routes import DeletePrompt
|
|
|
|
with app.test_request_context(
|
|
"/api/delete_prompt",
|
|
method="POST",
|
|
json={},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = DeletePrompt().post()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestUpdatePrompt:
|
|
pass
|
|
|
|
def test_returns_400_missing_fields(self, app):
|
|
from application.api.user.prompts.routes import UpdatePrompt
|
|
|
|
with app.test_request_context(
|
|
"/api/update_prompt",
|
|
method="POST",
|
|
json={"id": str(uuid.uuid4().hex[:24]), "name": "Updated"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user1"}
|
|
response = UpdatePrompt().post()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Happy-path tests using the ephemeral pg_conn fixture
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCreatePromptHappyPath:
|
|
def test_creates_prompt_returns_id(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import CreatePrompt
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "P1", "content": "c1"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "user-create"}
|
|
response = CreatePrompt().post()
|
|
|
|
assert response.status_code == 200
|
|
assert "id" in response.json
|
|
|
|
def test_create_error_returns_400(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import CreatePrompt
|
|
|
|
# Force repository error by closing the connection first
|
|
@contextmanager
|
|
def _broken():
|
|
raise RuntimeError("simulated db error")
|
|
yield # unreachable
|
|
|
|
with patch(
|
|
"application.api.user.prompts.routes.db_session", _broken
|
|
), app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "P", "content": "c"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "u1"}
|
|
response = CreatePrompt().post()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
class TestGetPromptsHappyPath:
|
|
def test_returns_builtin_plus_user_prompts(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import CreatePrompt, GetPrompts
|
|
|
|
user = "user-list"
|
|
# Seed two prompts via the same endpoint
|
|
for name in ("alpha", "beta"):
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": name, "content": f"content-{name}"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
CreatePrompt().post()
|
|
|
|
with _patch_db(pg_conn), app.test_request_context("/api/get_prompts"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = GetPrompts().get()
|
|
|
|
assert response.status_code == 200
|
|
names = [p["name"] for p in response.json]
|
|
# Three built-ins always present
|
|
assert "default" in names and "creative" in names and "strict" in names
|
|
assert "alpha" in names and "beta" in names
|
|
|
|
def test_get_error_returns_400(self, app):
|
|
from application.api.user.prompts.routes import GetPrompts
|
|
|
|
@contextmanager
|
|
def _broken():
|
|
raise RuntimeError("simulated db error")
|
|
yield
|
|
|
|
with patch(
|
|
"application.api.user.prompts.routes.db_readonly", _broken
|
|
), app.test_request_context("/api/get_prompts"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "u1"}
|
|
response = GetPrompts().get()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
class TestGetSinglePromptHappyPath:
|
|
def test_returns_private_prompt_content(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import (
|
|
CreatePrompt,
|
|
GetSinglePrompt,
|
|
)
|
|
|
|
user = "user-get1"
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "custom", "content": "hello world"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
created = CreatePrompt().post()
|
|
prompt_id = created.json["id"]
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
f"/api/get_single_prompt?id={prompt_id}"
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["content"] == "hello world"
|
|
|
|
def test_returns_404_for_unknown_prompt(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
bogus_id = str(uuid.uuid4())
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
f"/api/get_single_prompt?id={bogus_id}"
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "whoever"}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 404
|
|
|
|
def test_file_read_exception_returns_400(self, app):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
# Presets are composed in-process now, so the failure this covers is a
|
|
# repository error on the custom-prompt path.
|
|
with patch(
|
|
"application.api.user.prompts.routes.PromptsRepository",
|
|
side_effect=OSError("boom"),
|
|
), app.test_request_context("/api/get_single_prompt?id=some-custom-id"):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "u1"}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
class TestDeletePromptHappyPath:
|
|
def test_deletes_existing_prompt(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import (
|
|
CreatePrompt,
|
|
DeletePrompt,
|
|
GetSinglePrompt,
|
|
)
|
|
|
|
user = "user-del"
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "to-delete", "content": "bye"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
created = CreatePrompt().post()
|
|
prompt_id = created.json["id"]
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/delete_prompt",
|
|
method="POST",
|
|
json={"id": prompt_id},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = DeletePrompt().post()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["success"] is True
|
|
|
|
# Verify gone
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
f"/api/get_single_prompt?id={prompt_id}"
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
check = GetSinglePrompt().get()
|
|
assert check.status_code == 404
|
|
|
|
def test_delete_returns_401_unauthenticated(self, app):
|
|
from application.api.user.prompts.routes import DeletePrompt
|
|
|
|
with app.test_request_context(
|
|
"/api/delete_prompt",
|
|
method="POST",
|
|
json={"id": "something"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = None
|
|
response = DeletePrompt().post()
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_delete_error_returns_400(self, app):
|
|
from application.api.user.prompts.routes import DeletePrompt
|
|
|
|
@contextmanager
|
|
def _broken():
|
|
raise RuntimeError("boom")
|
|
yield
|
|
|
|
with patch(
|
|
"application.api.user.prompts.routes.db_session", _broken
|
|
), app.test_request_context(
|
|
"/api/delete_prompt",
|
|
method="POST",
|
|
json={"id": "pid"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "u1"}
|
|
response = DeletePrompt().post()
|
|
|
|
assert response.status_code == 400
|
|
|
|
|
|
class TestLegacyMongoIdResolution:
|
|
"""Pre-cutover prompt ids (Mongo ObjectIds) must resolve on all three
|
|
mutating endpoints once the prompts.legacy_mongo_id column is populated
|
|
by backfill. Previously the route short-circuited on non-UUID input via
|
|
``CAST(:id AS uuid)`` which raised and poisoned the transaction."""
|
|
|
|
LEGACY_ID = "507f1f77bcf86cd799439011"
|
|
|
|
def _seed_legacy(self, pg_conn, user: str, name: str, content: str):
|
|
from application.storage.db.repositories.prompts import PromptsRepository
|
|
|
|
return PromptsRepository(pg_conn).create(
|
|
user, name, content, legacy_mongo_id=self.LEGACY_ID,
|
|
)
|
|
|
|
def test_get_single_prompt_resolves_legacy_id(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import GetSinglePrompt
|
|
|
|
user = "legacy-user"
|
|
self._seed_legacy(pg_conn, user, "orig", "legacy-body")
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
f"/api/get_single_prompt?id={self.LEGACY_ID}"
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = GetSinglePrompt().get()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["content"] == "legacy-body"
|
|
|
|
def test_delete_prompt_resolves_legacy_id(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import DeletePrompt
|
|
from application.storage.db.repositories.prompts import PromptsRepository
|
|
|
|
user = "legacy-user-del"
|
|
self._seed_legacy(pg_conn, user, "to-delete", "x")
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/delete_prompt",
|
|
method="POST",
|
|
json={"id": self.LEGACY_ID},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = DeletePrompt().post()
|
|
|
|
assert response.status_code == 200
|
|
assert response.json["success"] is True
|
|
assert PromptsRepository(pg_conn).get_by_legacy_id(
|
|
self.LEGACY_ID, user,
|
|
) is None
|
|
|
|
def test_update_prompt_resolves_legacy_id(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import UpdatePrompt
|
|
from application.storage.db.repositories.prompts import PromptsRepository
|
|
|
|
user = "legacy-user-upd"
|
|
self._seed_legacy(pg_conn, user, "old-name", "old-content")
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/update_prompt",
|
|
method="POST",
|
|
json={"id": self.LEGACY_ID, "name": "new-name", "content": "new-content"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = UpdatePrompt().post()
|
|
|
|
assert response.status_code == 200
|
|
fetched = PromptsRepository(pg_conn).get_by_legacy_id(self.LEGACY_ID, user)
|
|
assert fetched["name"] == "new-name"
|
|
assert fetched["content"] == "new-content"
|
|
|
|
|
|
class TestUpdatePromptHappyPath:
|
|
def test_updates_prompt(self, app, pg_conn):
|
|
from application.api.user.prompts.routes import (
|
|
CreatePrompt,
|
|
GetSinglePrompt,
|
|
UpdatePrompt,
|
|
)
|
|
|
|
user = "user-upd"
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/create_prompt",
|
|
method="POST",
|
|
json={"name": "orig", "content": "v1"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
created = CreatePrompt().post()
|
|
prompt_id = created.json["id"]
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
"/api/update_prompt",
|
|
method="POST",
|
|
json={"id": prompt_id, "name": "renamed", "content": "v2"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
response = UpdatePrompt().post()
|
|
assert response.status_code == 200
|
|
|
|
with _patch_db(pg_conn), app.test_request_context(
|
|
f"/api/get_single_prompt?id={prompt_id}"
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": user}
|
|
check = GetSinglePrompt().get()
|
|
assert check.status_code == 200
|
|
assert check.json["content"] == "v2"
|
|
|
|
def test_update_returns_401_unauthenticated(self, app):
|
|
from application.api.user.prompts.routes import UpdatePrompt
|
|
|
|
with app.test_request_context(
|
|
"/api/update_prompt",
|
|
method="POST",
|
|
json={"id": "x", "name": "n", "content": "c"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = None
|
|
response = UpdatePrompt().post()
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_update_error_returns_400(self, app):
|
|
from application.api.user.prompts.routes import UpdatePrompt
|
|
|
|
@contextmanager
|
|
def _broken():
|
|
raise RuntimeError("boom")
|
|
yield
|
|
|
|
with patch(
|
|
"application.api.user.prompts.routes.db_session", _broken
|
|
), app.test_request_context(
|
|
"/api/update_prompt",
|
|
method="POST",
|
|
json={"id": "x", "name": "n", "content": "c"},
|
|
):
|
|
from flask import request
|
|
|
|
request.decoded_token = {"sub": "u1"}
|
|
response = UpdatePrompt().post()
|
|
|
|
assert response.status_code == 400
|