mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 08:13:02 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
437 lines
16 KiB
Python
437 lines
16 KiB
Python
"""Tests for the journaled execute path on ToolExecutor.
|
|
|
|
Each tool call inserts a ``tool_call_attempts`` row and flips it
|
|
``proposed → executed`` (or ``→ failed``). With a ``message_id`` it
|
|
stays ``executed`` for the finalize path to confirm; without one
|
|
(``save_conversation=False``) it goes straight to ``confirmed``.
|
|
"""
|
|
|
|
from contextlib import contextmanager
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
|
|
from docsgpt.agents.tool_executor import ToolExecutor
|
|
|
|
|
|
@contextmanager
|
|
def _yield_pg(conn):
|
|
"""Adapter so the executor's ``db_session()`` writes land on ``pg_conn``."""
|
|
|
|
@contextmanager
|
|
def _yield():
|
|
yield conn
|
|
|
|
return _yield
|
|
|
|
|
|
def _patch_db(monkeypatch, pg_conn):
|
|
"""Patch all ``db_session`` entry points used by the executor and tools.
|
|
|
|
Each module imports ``db_session`` / ``db_readonly`` by name so each
|
|
module-level binding has to be replaced individually.
|
|
"""
|
|
|
|
@contextmanager
|
|
def _use_pg():
|
|
yield pg_conn
|
|
|
|
targets = (
|
|
"docsgpt.agents.tool_executor",
|
|
"docsgpt.agents.tools.notes",
|
|
"docsgpt.agents.tools.todo_list",
|
|
"docsgpt.storage.db.session",
|
|
)
|
|
for module in targets:
|
|
monkeypatch.setattr(f"{module}.db_session", _use_pg, raising=False)
|
|
monkeypatch.setattr(f"{module}.db_readonly", _use_pg, raising=False)
|
|
|
|
|
|
def _drain(gen):
|
|
"""Exhaust a generator, returning ``(events, return_value)``."""
|
|
events = []
|
|
while True:
|
|
try:
|
|
events.append(next(gen))
|
|
except StopIteration as exc:
|
|
return events, exc.value
|
|
|
|
|
|
def _select_attempt(pg_conn, call_id):
|
|
row = pg_conn.execute(
|
|
text("SELECT * FROM tool_call_attempts WHERE call_id = :cid"),
|
|
{"cid": call_id},
|
|
).fetchone()
|
|
return row._mapping if row is not None else None
|
|
|
|
|
|
def _make_call(name="test_action_t1", call_id="c1"):
|
|
call = Mock()
|
|
call.name = name
|
|
call.id = call_id
|
|
call.arguments = "{}"
|
|
return call
|
|
|
|
|
|
_TOOLS_DICT = {
|
|
"t1": {
|
|
"id": "00000000-0000-0000-0000-000000000001",
|
|
"name": "test_tool",
|
|
"config": {"key": "val"},
|
|
"actions": [
|
|
{"name": "test_action", "description": "T", "parameters": {"properties": {}}},
|
|
],
|
|
}
|
|
}
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExecuteJournaling:
|
|
def test_no_message_id_proposed_then_confirmed(
|
|
self, pg_conn, mock_tool_manager, monkeypatch
|
|
):
|
|
"""No reserved message (``save_conversation=False``) → row lands ``confirmed``, not ``executed``."""
|
|
executor = ToolExecutor(user="u")
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolActionParser",
|
|
lambda _cls, **kw: Mock(
|
|
parse_args=Mock(return_value=("t1", "test_action", {"q": "v"}))
|
|
),
|
|
)
|
|
_patch_db(monkeypatch, pg_conn)
|
|
|
|
events, result = _drain(executor.execute(_TOOLS_DICT, _make_call(), "MockLLM"))
|
|
assert result[0] == "Tool result"
|
|
|
|
row = _select_attempt(pg_conn, "c1")
|
|
assert row is not None
|
|
assert row["status"] == "confirmed"
|
|
assert row["tool_name"] == "test_tool"
|
|
assert row["action_name"] == "test_action"
|
|
assert row["arguments"] == {"q": "v"}
|
|
# Result is wrapped so a future ``artifact_id`` can ride alongside.
|
|
assert row["result"] == {"result": "Tool result"}
|
|
assert row["error"] is None
|
|
assert row["message_id"] is None
|
|
|
|
def test_executor_message_id_is_persisted_on_executed_row(
|
|
self, pg_conn, mock_tool_manager, monkeypatch
|
|
):
|
|
"""The executor's message_id is carried onto the journal row, which stays ``executed``."""
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
|
|
# FK constraint: message_id must reference a real row.
|
|
repo = ConversationsRepository(pg_conn)
|
|
conv = repo.create("u-mid", "msg-id-test")
|
|
msg = repo.reserve_message(
|
|
str(conv["id"]),
|
|
prompt="q?",
|
|
placeholder_response="...",
|
|
request_id="req-mid-1",
|
|
status="pending",
|
|
)
|
|
message_uuid = str(msg["id"])
|
|
|
|
executor = ToolExecutor(user="u")
|
|
executor.message_id = message_uuid
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolActionParser",
|
|
lambda _cls, **kw: Mock(
|
|
parse_args=Mock(return_value=("t1", "test_action", {}))
|
|
),
|
|
)
|
|
_patch_db(monkeypatch, pg_conn)
|
|
|
|
_drain(executor.execute(_TOOLS_DICT, _make_call(call_id="cm1"), "MockLLM"))
|
|
|
|
# The durability key is namespaced by message_id so providers that reuse
|
|
# deterministic call ids across turns don't collide on the table-wide PK;
|
|
# the raw id alone therefore matches nothing.
|
|
assert _select_attempt(pg_conn, "cm1") is None
|
|
row = _select_attempt(pg_conn, f"{message_uuid}:cm1")
|
|
assert row is not None
|
|
assert row["status"] == "executed"
|
|
assert str(row["message_id"]) == message_uuid
|
|
|
|
def test_reused_call_id_across_turns_does_not_collide(
|
|
self, pg_conn, mock_tool_manager, monkeypatch
|
|
):
|
|
"""Regression for the tool_call_attempts PK collision (#9): a provider
|
|
that reuses a deterministic call id (``functions.create_artifact:0``)
|
|
across turns must journal a distinct row per turn, not silently drop the
|
|
later one on the table-wide ``call_id`` primary key."""
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
|
|
repo = ConversationsRepository(pg_conn)
|
|
conv = repo.create("u-collide", "collision-test")
|
|
msg1 = repo.reserve_message(
|
|
str(conv["id"]), prompt="q1", placeholder_response="...",
|
|
request_id="req-c1", status="pending",
|
|
)
|
|
msg2 = repo.reserve_message(
|
|
str(conv["id"]), prompt="q2", placeholder_response="...",
|
|
request_id="req-c2", status="pending",
|
|
)
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolActionParser",
|
|
lambda _cls, **kw: Mock(
|
|
parse_args=Mock(return_value=("t1", "test_action", {}))
|
|
),
|
|
)
|
|
_patch_db(monkeypatch, pg_conn)
|
|
|
|
reused = "functions.create_artifact:0"
|
|
for msg in (msg1, msg2):
|
|
executor = ToolExecutor(user="u-collide")
|
|
executor.message_id = str(msg["id"])
|
|
_drain(
|
|
executor.execute(_TOOLS_DICT, _make_call(call_id=reused), "MockLLM")
|
|
)
|
|
|
|
# Both turns journaled their own row; before the fix the second call hit
|
|
# ON CONFLICT DO NOTHING on the shared PK and was dropped.
|
|
assert _select_attempt(pg_conn, reused) is None
|
|
r1 = _select_attempt(pg_conn, f"{msg1['id']}:{reused}")
|
|
r2 = _select_attempt(pg_conn, f"{msg2['id']}:{reused}")
|
|
assert r1 is not None and str(r1["message_id"]) == str(msg1["id"])
|
|
assert r2 is not None and str(r2["message_id"]) == str(msg2["id"])
|
|
|
|
def test_tool_raises_marks_failed_and_reraises(
|
|
self, pg_conn, mock_tool_manager, monkeypatch
|
|
):
|
|
executor = ToolExecutor(user="u")
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolActionParser",
|
|
lambda _cls, **kw: Mock(
|
|
parse_args=Mock(return_value=("t1", "test_action", {}))
|
|
),
|
|
)
|
|
_patch_db(monkeypatch, pg_conn)
|
|
mock_tool_manager.load_tool.return_value.execute_action.side_effect = (
|
|
RuntimeError("boom")
|
|
)
|
|
|
|
gen = executor.execute(_TOOLS_DICT, _make_call(call_id="c2"), "MockLLM")
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
_drain(gen)
|
|
|
|
row = _select_attempt(pg_conn, "c2")
|
|
assert row is not None
|
|
assert row["status"] == "failed"
|
|
assert row["error"] == "boom"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestRepository:
|
|
def test_proposed_then_confirmed_when_no_message(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
assert repo.record_proposed("c-x", "tool", "act", {"a": 1}) is True
|
|
# Duplicate insert is a no-op; original row stays put.
|
|
assert repo.record_proposed("c-x", "tool", "act", {"a": 1}) is False
|
|
row = _select_attempt(pg_conn, "c-x")
|
|
assert row["status"] == "proposed"
|
|
|
|
assert repo.mark_executed("c-x", {"out": "ok"}) is True
|
|
row = _select_attempt(pg_conn, "c-x")
|
|
assert row["status"] == "confirmed"
|
|
assert row["message_id"] is None
|
|
assert row["result"] == {"result": {"out": "ok"}}
|
|
|
|
def test_mark_executed_with_message_stays_executed(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
# FK constraint: message_id must reference a real row.
|
|
conv_repo = ConversationsRepository(pg_conn)
|
|
conv = conv_repo.create("u-repo", "repo-msg-test")
|
|
msg = conv_repo.reserve_message(
|
|
str(conv["id"]),
|
|
prompt="q?",
|
|
placeholder_response="...",
|
|
request_id="req-repo-1",
|
|
status="pending",
|
|
)
|
|
message_uuid = str(msg["id"])
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("c-m", "tool", "act", {})
|
|
assert (
|
|
repo.mark_executed("c-m", {"out": "ok"}, message_id=message_uuid) is True
|
|
)
|
|
row = _select_attempt(pg_conn, "c-m")
|
|
assert row["status"] == "executed"
|
|
assert str(row["message_id"]) == message_uuid
|
|
|
|
def test_upsert_executed_without_message_confirms(self, pg_conn):
|
|
"""``upsert_executed`` (DB-outage fallback) with no ``message_id`` lands ``confirmed``."""
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.upsert_executed("c-up", "tool", "act", {"a": 1}, {"out": "ok"})
|
|
row = _select_attempt(pg_conn, "c-up")
|
|
assert row["status"] == "confirmed"
|
|
assert row["message_id"] is None
|
|
assert row["result"] == {"result": {"out": "ok"}}
|
|
|
|
def test_mark_failed_sets_error(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("c-y", "tool", "act", {})
|
|
assert repo.mark_failed("c-y", "kaboom") is True
|
|
row = _select_attempt(pg_conn, "c-y")
|
|
assert row["status"] == "failed"
|
|
assert row["error"] == "kaboom"
|
|
|
|
def test_mark_failed_leaves_executed_row_untouched(self, pg_conn):
|
|
"""A late error for a reused ``call_id`` ("call_0"-style) must
|
|
not flip an already-executed row (see ``mark_failed``)."""
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("c-dup", "tool", "act", {})
|
|
# No message_id ⇒ lands straight in 'confirmed' (still non-proposed).
|
|
assert repo.mark_executed("c-dup", {"out": "ok"}) is True
|
|
# Late error path hits the same id — the guard rejects the update.
|
|
assert repo.mark_failed("c-dup", "boom") is False
|
|
row = _select_attempt(pg_conn, "c-dup")
|
|
assert row["status"] == "confirmed"
|
|
assert row["error"] is None
|
|
|
|
def test_mark_executed_guarded_to_proposed(self, pg_conn):
|
|
"""A reused ``call_id`` must not flip an already-terminal row back
|
|
to executed (the status guard mirrors ``mark_failed``)."""
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("c-me", "tool", "act", {})
|
|
assert repo.mark_failed("c-me", "boom") is True
|
|
# Late executed emission for the same id: no proposed row remains.
|
|
assert repo.mark_executed("c-me", {"out": "ok"}) is False
|
|
row = _select_attempt(pg_conn, "c-me")
|
|
assert row["status"] == "failed"
|
|
assert row["result"] is None
|
|
|
|
def test_mark_executed_scoped_to_owner(self, pg_conn):
|
|
"""A colliding ``call_id`` from another tenant can't flip — or
|
|
read its result into — this user's proposed row."""
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("call_0", "tool", "act", {}, user_id="victim")
|
|
# Attacker's request reuses the id and tries to land its result.
|
|
assert (
|
|
repo.mark_executed("call_0", {"out": "leak"}, user_id="attacker")
|
|
is False
|
|
)
|
|
row = _select_attempt(pg_conn, "call_0")
|
|
assert row["status"] == "proposed"
|
|
assert row["result"] is None
|
|
assert row["user_id"] == "victim"
|
|
|
|
def test_mark_failed_scoped_to_owner(self, pg_conn):
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("call_0", "tool", "act", {}, user_id="victim")
|
|
assert repo.mark_failed("call_0", "boom", user_id="attacker") is False
|
|
row = _select_attempt(pg_conn, "call_0")
|
|
assert row["status"] == "proposed"
|
|
|
|
def test_upsert_executed_stamps_attribution(self, pg_conn):
|
|
"""The DB-outage fallback row must carry user/agent so it stays
|
|
visible to per-user / per-agent analytics (was born unattributed)."""
|
|
import uuid as _uuid
|
|
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
agent_id = str(_uuid.uuid4())
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.upsert_executed(
|
|
"c-up2", "tool", "act", {}, {"out": "ok"},
|
|
user_id="u-up", agent_id=agent_id,
|
|
)
|
|
row = _select_attempt(pg_conn, "c-up2")
|
|
assert row["user_id"] == "u-up"
|
|
assert str(row["agent_id"]) == agent_id
|
|
|
|
def test_upsert_executed_wont_clobber_other_tenant(self, pg_conn):
|
|
"""A colliding fallback upsert must not upgrade another tenant's
|
|
proposed row or overwrite its result."""
|
|
from docsgpt.storage.db.repositories.tool_call_attempts import (
|
|
ToolCallAttemptsRepository,
|
|
)
|
|
|
|
repo = ToolCallAttemptsRepository(pg_conn)
|
|
repo.record_proposed("call_0", "tool", "act", {}, user_id="victim")
|
|
repo.upsert_executed(
|
|
"call_0", "tool", "act", {}, {"out": "leak"}, user_id="attacker"
|
|
)
|
|
row = _select_attempt(pg_conn, "call_0")
|
|
assert row["status"] == "proposed"
|
|
assert row["result"] is None
|
|
assert row["user_id"] == "victim"
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestDefaultToolJournaling:
|
|
"""A default tool's synthetic id round-trips through execute/journal."""
|
|
|
|
def test_synthetic_tool_id_is_journaled(
|
|
self, pg_conn, mock_tool_manager, monkeypatch
|
|
):
|
|
from docsgpt.agents.default_tools import synthesize_default_tool
|
|
|
|
memory_row = synthesize_default_tool("memory")
|
|
assert memory_row is not None
|
|
tools_dict = {memory_row["id"]: memory_row}
|
|
|
|
executor = ToolExecutor(user="u")
|
|
monkeypatch.setattr(
|
|
"docsgpt.agents.tool_executor.ToolActionParser",
|
|
lambda _cls, **kw: Mock(
|
|
parse_args=Mock(
|
|
return_value=(memory_row["id"], "memory_view", {"path": "/"})
|
|
)
|
|
),
|
|
)
|
|
_patch_db(monkeypatch, pg_conn)
|
|
|
|
events, result = _drain(
|
|
executor.execute(tools_dict, _make_call(call_id="c-def"), "MockLLM")
|
|
)
|
|
assert result[0] == "Tool result"
|
|
|
|
row = _select_attempt(pg_conn, "c-def")
|
|
assert row is not None
|
|
assert row["status"] == "confirmed"
|
|
assert row["tool_name"] == "memory"
|
|
assert str(row["tool_id"]) == memory_row["id"]
|