101 lines
4.9 KiB
Python
101 lines
4.9 KiB
Python
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from app.gateway import assistant_file_ingest as ingest
|
|
from app.gateway.routers import assistant_knowledge as router
|
|
from deerflow.knowledge.schemas import NoteDraft
|
|
from deerflow.persistence.assistant_knowledge import AssistantKnowledgeRepository
|
|
from deerflow.persistence.base import Base
|
|
|
|
|
|
def test_same_topic_across_document_parts_preserves_all_knowledge():
|
|
package = ingest.notes_package([
|
|
{"title": "设备维护", "content_md": "每周清理过滤网", "tags": ["运维"]},
|
|
{"title": "设备 维护", "content_md": "每月检查电机温度", "tags": ["设备"]},
|
|
], {"digest": "document", "filename": "维护指南.md"})
|
|
assert len(package["wiki_pages"]) == 1
|
|
assert "每周清理过滤网" in package["wiki_pages"][0]["content"]
|
|
assert "每月检查电机温度" in package["wiki_pages"][0]["content"]
|
|
assert package["wiki_pages"][0]["tags"] == ["运维", "设备"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_standalone_upload_vectors_retry_and_dedup(tmp_path, monkeypatch):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
store = AssistantKnowledgeRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
base = await store.create_base(name="独立文档库", created_by="admin")
|
|
app = FastAPI()
|
|
app.include_router(router.router)
|
|
app.state.assistant_knowledge_store = store
|
|
app.state.llmwiki_embedding = None
|
|
monkeypatch.setattr(ingest, "runtime_home", lambda: tmp_path)
|
|
extraction_calls = []
|
|
|
|
async def actor(*args, **kwargs):
|
|
return "admin"
|
|
|
|
async def extract(text, **kwargs):
|
|
extraction_calls.append(text)
|
|
return [NoteDraft(title="冷却系统", summary="摘要占位", content_md="# 冷却系统\n风扇的额定转速为每分钟 1200 转。", tags=["设备"], entities=[{"name": "风扇", "type": "product"}])]
|
|
|
|
class Encoder:
|
|
fingerprint = "local-test-3"
|
|
dimensions = 3
|
|
config = SimpleNamespace(model="offline-test")
|
|
|
|
async def embed_texts(self, texts):
|
|
return [[1.0, 0.0, 0.0] for _ in texts]
|
|
|
|
async def embed_query(self, text):
|
|
return [1.0, 0.0, 0.0]
|
|
|
|
monkeypatch.setattr(router, "_actor", actor)
|
|
monkeypatch.setattr(ingest, "extract_document_llm_multi", extract)
|
|
|
|
async def finish():
|
|
if ingest._tasks:
|
|
await asyncio.gather(*list(ingest._tasks.values()))
|
|
|
|
try:
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
|
payload = {"assistant_base_id": base["id"]}
|
|
response = await client.post("/api/assistant-knowledge/imports/files", data=payload, files={"file": ("../../设备.md", "风扇额定转速为 1200 转。".encode())})
|
|
assert response.status_code == 202, response.text
|
|
job_id = response.json()["id"]
|
|
await finish()
|
|
assert (await store.get_job(job_id))["status"] == "failed"
|
|
assert (ingest.file_job_dir(job_id) / "wiki.json").exists()
|
|
app.state.llmwiki_embedding = Encoder()
|
|
retry = await client.post(f"/api/assistant-knowledge/import-jobs/{job_id}/retry-file")
|
|
assert retry.status_code == 202
|
|
await finish()
|
|
job = await store.get_job(job_id)
|
|
assert job["status"] == "completed", job
|
|
assert job["counts"]["vectors"] == 1
|
|
assert len(extraction_calls) == 1
|
|
results = await store.search_pages(base_ids=[base["id"]], query="额定转速", embedding_client=app.state.llmwiki_embedding)
|
|
assert results and results[0]["retrieval_mode"] == "vector"
|
|
assert "1200" in results[0]["matched_section"]["content"]
|
|
assert "摘要占位" not in results[0]["matched_section"]["content"]
|
|
global_base = await store.get_base()
|
|
assert len(await store.list_pages(base_id=global_base["id"])) == 1
|
|
again = await client.post("/api/assistant-knowledge/imports/files", data=payload, files={"file": ("设备.md", "风扇额定转速为 1200 转。".encode())})
|
|
assert again.status_code == 202
|
|
await finish()
|
|
assert len(await store.list_sources(base_id=base["id"])) == 1
|
|
assert len(await store.list_pages(base_id=global_base["id"])) == 1
|
|
bad = await client.post("/api/assistant-knowledge/imports/files", data=payload, files={"file": ("script.py", b"print('no')")})
|
|
assert bad.status_code == 415
|
|
empty = await client.post("/api/assistant-knowledge/imports/files", data=payload, files={"file": ("empty.txt", b"")})
|
|
assert empty.status_code == 422
|
|
finally:
|
|
await ingest.stop_file_jobs()
|
|
await engine.dispose()
|