fix(avatar): index knowledge documents asynchronously
This commit is contained in:
@@ -8,6 +8,7 @@ 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)
|
||||
@@ -31,14 +32,14 @@ def test_doc_payload_reports_whether_the_persisted_file_exists(tmp_path: Path):
|
||||
assert _doc_payload(doc)["filePresent"] is True
|
||||
|
||||
|
||||
def test_upload_marks_vectorization_failure_instead_of_staying_processing(
|
||||
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.embeddings.embed", side_effect=RuntimeError("provider unavailable")),
|
||||
patch("routers.knowledge.knowledge_vectorizer.enqueue") as enqueue,
|
||||
):
|
||||
response = client.post(
|
||||
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
||||
@@ -47,14 +48,15 @@ def test_upload_marks_vectorization_failure_instead_of_staying_processing(
|
||||
)
|
||||
|
||||
payload = response.json()["data"]
|
||||
assert payload["status"] == "failed"
|
||||
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 == "failed"
|
||||
assert stored.status == "parsing"
|
||||
assert db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == stored.id).count() == 0
|
||||
db.delete(stored)
|
||||
db.commit()
|
||||
@@ -62,14 +64,14 @@ def test_upload_marks_vectorization_failure_instead_of_staying_processing(
|
||||
db.close()
|
||||
|
||||
|
||||
def test_markdown_upload_commits_ready_document_and_chunks_together(
|
||||
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.embeddings.embed", return_value=[[1.0, 0.0]]),
|
||||
patch("routers.knowledge.knowledge_vectorizer.enqueue"),
|
||||
):
|
||||
response = client.post(
|
||||
f"/api/avatar/{context['avatar'].id}/knowledge/docs",
|
||||
@@ -78,14 +80,19 @@ def test_markdown_upload_commits_ready_document_and_chunks_together(
|
||||
)
|
||||
|
||||
payload = response.json()["data"]
|
||||
assert payload["status"] == "ready"
|
||||
assert payload["vectorized"] is True
|
||||
assert payload["chunkCount"] == 1
|
||||
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)
|
||||
@@ -94,6 +101,87 @@ def test_markdown_upload_commits_ready_document_and_chunks_together(
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user