297 lines
10 KiB
Python
297 lines
10 KiB
Python
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from database import SessionLocal
|
|
from main import app
|
|
from models import Avatar, KnowledgeChunk, KnowledgeDoc, QAPair
|
|
from routers.knowledge import _doc_payload
|
|
from services.knowledge_vectorizer import knowledge_vectorizer
|
|
|
|
|
|
client = TestClient(app)
|
|
|
|
|
|
def test_doc_payload_reports_whether_the_persisted_file_exists(tmp_path: Path):
|
|
avatar_id = "avatar-1"
|
|
stored_name = "knowledge.md"
|
|
doc = SimpleNamespace(
|
|
avatar_id=avatar_id,
|
|
file_url=f"/api/files/{avatar_id}/{stored_name}",
|
|
to_dict=lambda: {"id": "doc-1", "fileUrl": f"/api/files/{avatar_id}/{stored_name}"},
|
|
)
|
|
stored_dir = tmp_path / avatar_id
|
|
stored_dir.mkdir()
|
|
stored_file = stored_dir / stored_name
|
|
|
|
with patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)):
|
|
assert _doc_payload(doc)["filePresent"] is False
|
|
stored_file.write_text("knowledge", encoding="utf-8")
|
|
assert _doc_payload(doc)["filePresent"] is True
|
|
|
|
|
|
def test_upload_returns_before_background_vectorization(
|
|
tmp_path: Path,
|
|
authorization_context,
|
|
):
|
|
context = authorization_context
|
|
with (
|
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
|
patch("routers.knowledge.knowledge_vectorizer.enqueue") as enqueue,
|
|
):
|
|
response = client.post(
|
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
|
headers=context["owner_headers"],
|
|
files={"file": ("knowledge.md", b"# Knowledge\n\nTest content", "text/markdown")},
|
|
)
|
|
|
|
payload = response.json()["data"]
|
|
assert payload["status"] == "parsing"
|
|
assert payload["vectorized"] is False
|
|
assert payload["chunkCount"] == 0
|
|
enqueue.assert_called_once_with(payload["id"])
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one()
|
|
assert stored.status == "parsing"
|
|
assert db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).count() == 0
|
|
db.delete(stored)
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def test_upload_rejects_oversize_file_before_queuing_indexing(
|
|
tmp_path: Path,
|
|
authorization_context,
|
|
):
|
|
context = authorization_context
|
|
with (
|
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
|
patch("routers.knowledge.MAX_UPLOAD_BYTES", 4),
|
|
patch("routers.knowledge.knowledge_vectorizer.enqueue") as enqueue,
|
|
):
|
|
response = client.post(
|
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
|
headers=context["owner_headers"],
|
|
files={"file": ("oversize.md", b"12345", "text/markdown")},
|
|
)
|
|
|
|
payload = response.json()
|
|
assert payload["code"] == 400
|
|
assert payload["message"] == "文件不能超过 50MB"
|
|
enqueue.assert_not_called()
|
|
assert not list((tmp_path / context["avatar"].id).glob("*"))
|
|
|
|
|
|
def test_background_vectorizer_commits_ready_document_and_chunks_together(
|
|
tmp_path: Path,
|
|
authorization_context,
|
|
):
|
|
context = authorization_context
|
|
with (
|
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
|
patch("routers.knowledge.knowledge_vectorizer.enqueue"),
|
|
):
|
|
response = client.post(
|
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
|
headers=context["owner_headers"],
|
|
files={"file": ("knowledge.md", b"# Knowledge\n\nTest content", "text/markdown")},
|
|
)
|
|
|
|
payload = response.json()["data"]
|
|
assert payload["status"] == "parsing"
|
|
with (
|
|
patch("services.knowledge_vectorizer.UPLOAD_DIR", str(tmp_path)),
|
|
patch("services.knowledge_vectorizer.embeddings.embed", return_value=[[1.0, 0.0]]),
|
|
):
|
|
knowledge_vectorizer.vectorize_document(payload["id"])
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one()
|
|
assert stored.status == "ready"
|
|
assert stored.vectorized is True
|
|
assert stored.chunk_count == 1
|
|
assert db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).count() == 1
|
|
db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).delete()
|
|
db.delete(stored)
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def test_background_vectorizer_keeps_failure_reason_for_retry(
|
|
tmp_path: Path,
|
|
authorization_context,
|
|
):
|
|
context = authorization_context
|
|
with (
|
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
|
patch("routers.knowledge.knowledge_vectorizer.enqueue"),
|
|
):
|
|
response = client.post(
|
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
|
headers=context["owner_headers"],
|
|
files={"file": ("knowledge.md", b"# Knowledge\n\nTest content", "text/markdown")},
|
|
)
|
|
|
|
payload = response.json()["data"]
|
|
with (
|
|
patch("services.knowledge_vectorizer.UPLOAD_DIR", str(tmp_path)),
|
|
patch("services.knowledge_vectorizer.embeddings.embed", side_effect=RuntimeError("provider unavailable")),
|
|
):
|
|
knowledge_vectorizer.vectorize_document(payload["id"])
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
stored = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == payload["id"]).one()
|
|
assert stored.status == "failed"
|
|
assert stored.error_message == "provider unavailable"
|
|
assert db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).count() == 0
|
|
db.delete(stored)
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def test_retry_queues_a_failed_document_again(
|
|
tmp_path: Path,
|
|
authorization_context,
|
|
):
|
|
context = authorization_context
|
|
document_id = f"retry-doc-{context['suffix']}"
|
|
avatar_dir = tmp_path / context["avatar"].id
|
|
avatar_dir.mkdir()
|
|
(avatar_dir / "retry.md").write_text("retry content", encoding="utf-8")
|
|
db = SessionLocal()
|
|
try:
|
|
db.add(
|
|
KnowledgeDoc(
|
|
id=document_id,
|
|
avatar_id=context["avatar"].id,
|
|
filename="retry.md",
|
|
file_type="md",
|
|
file_url=f"/api/files/{context['avatar'].id}/retry.md",
|
|
status="failed",
|
|
error_message="provider unavailable",
|
|
)
|
|
)
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
with (
|
|
patch("routers.knowledge.UPLOAD_DIR", str(tmp_path)),
|
|
patch("routers.knowledge.knowledge_vectorizer.enqueue") as enqueue,
|
|
):
|
|
response = client.post(
|
|
f"/api/avatar/{context['avatar'].id}/knowledge/docs/{document_id}/retry",
|
|
headers=context["owner_headers"],
|
|
)
|
|
|
|
payload = response.json()["data"]
|
|
assert payload["status"] == "parsing"
|
|
assert payload["errorMessage"] == ""
|
|
enqueue.assert_called_once_with(document_id)
|
|
db = SessionLocal()
|
|
try:
|
|
db.query(KnowledgeDoc).filter(KnowledgeDoc.id == document_id).delete()
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def test_each_avatar_has_an_independent_document_and_qa_scope(authorization_context):
|
|
context = authorization_context
|
|
first_avatar_id = context["avatar"].id
|
|
second_avatar_id = f"knowledge-second-{context['suffix']}"
|
|
first_doc_id = f"knowledge-first-doc-{context['suffix']}"
|
|
second_doc_id = f"knowledge-second-doc-{context['suffix']}"
|
|
first_qa_id = f"knowledge-first-qa-{context['suffix']}"
|
|
second_qa_id = f"knowledge-second-qa-{context['suffix']}"
|
|
|
|
db = SessionLocal()
|
|
try:
|
|
db.add_all(
|
|
[
|
|
Avatar(
|
|
id=second_avatar_id,
|
|
owner_id=context["owner"].huihui_user_id,
|
|
name="独立知识库分身",
|
|
status="active",
|
|
config={},
|
|
),
|
|
KnowledgeDoc(
|
|
id=first_doc_id,
|
|
avatar_id=first_avatar_id,
|
|
filename="first.md",
|
|
status="ready",
|
|
vectorized=True,
|
|
),
|
|
KnowledgeDoc(
|
|
id=second_doc_id,
|
|
avatar_id=second_avatar_id,
|
|
filename="second.md",
|
|
status="ready",
|
|
vectorized=True,
|
|
),
|
|
QAPair(
|
|
id=first_qa_id,
|
|
avatar_id=first_avatar_id,
|
|
question="第一个分身问题",
|
|
answer="第一个分身答案",
|
|
),
|
|
QAPair(
|
|
id=second_qa_id,
|
|
avatar_id=second_avatar_id,
|
|
question="第二个分身问题",
|
|
answer="第二个分身答案",
|
|
),
|
|
]
|
|
)
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
try:
|
|
first_docs = client.get(
|
|
f"/api/avatar/{first_avatar_id}/knowledge/docs",
|
|
headers=context["owner_headers"],
|
|
).json()["data"]
|
|
second_docs = client.get(
|
|
f"/api/avatar/{second_avatar_id}/knowledge/docs",
|
|
headers=context["owner_headers"],
|
|
).json()["data"]
|
|
first_qa = client.get(
|
|
f"/api/avatar/{first_avatar_id}/knowledge/qa",
|
|
headers=context["owner_headers"],
|
|
).json()["data"]
|
|
second_qa = client.get(
|
|
f"/api/avatar/{second_avatar_id}/knowledge/qa",
|
|
headers=context["owner_headers"],
|
|
).json()["data"]
|
|
|
|
assert [item["id"] for item in first_docs if item["id"] == first_doc_id] == [first_doc_id]
|
|
assert second_doc_id not in {item["id"] for item in first_docs}
|
|
assert [item["id"] for item in second_docs] == [second_doc_id]
|
|
assert first_qa_id in {item["id"] for item in first_qa}
|
|
assert second_qa_id not in {item["id"] for item in first_qa}
|
|
assert [item["id"] for item in second_qa] == [second_qa_id]
|
|
finally:
|
|
db = SessionLocal()
|
|
try:
|
|
db.query(QAPair).filter(QAPair.id.in_([first_qa_id, second_qa_id])).delete(
|
|
synchronize_session=False
|
|
)
|
|
db.query(KnowledgeDoc).filter(
|
|
KnowledgeDoc.id.in_([first_doc_id, second_doc_id])
|
|
).delete(synchronize_session=False)
|
|
db.query(Avatar).filter(Avatar.id == second_avatar_id).delete()
|
|
db.commit()
|
|
finally:
|
|
db.close()
|