feat(deepdata): 引入修改密码、文件预览组件与清洗算子增强，deepflow 改名为 deep-flow
本次同步 09-12 至 09-28 的代码变更，主要内容： 新增功能 - 修改密码：后端新增 ChangePasswordRequest DTO 与UserController/UserService 对应接口，前端新增 ChangePasswordDialog 并在 Header 接入入口 - 文件预览组件重构：file-preview 目录拆分出 VideoPreview（支持 blobUrl 直传）、 ImagePreview、AudioPreview，index.tsx 统一调度各类型预览器 - 新增 my_ops 算子模块：markitdown_parser、mineru_pdf_parser、text_quality_evaluator - 新增 tools/minio-file-server MinIO 文件服务 功能增强 - 数据清洗：cleaning_task_service/scheduler/routes 重构， data_juicer_executor 与 wrapper 增强，OperatorTable/LogsTable/TaskDetail 同步调整 - 数据标注：annotation sync.py、project.py 接口与schema 更新 - 数据采集、算子市场、首页与菜单路由相关页面调整 - DataJuicer 算子接入（useOperatorOperations、dj_ops_schema） 结构调整 - deepflow/ 更名为 deep-flow/（419 个文件重命名） - 新增 helm 部署编排、GitHub Actions 之外的自建镜像构建脚本 未纳入本次提交 - scripts-backup/、runtime/python-executor-backup/、components_backup/（备份目录） - *.backup*、__pycache__ 产物 - zh_core_web_sm-3.8.0-py3-none-any.whl（46MB 构建产物） - DataX/（内嵌独立仓库，远程亦未收录其源码）
This commit is contained in:
沃巴蒂
@@ -0,0 +1,154 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.rag.dify import DifyProvider, parse_uri
|
||||
|
||||
|
||||
# Dummy classes to mock dependencies
|
||||
class DummyResource:
|
||||
def __init__(self, uri, title="", description=""):
|
||||
self.uri = uri
|
||||
self.title = title
|
||||
self.description = description
|
||||
|
||||
|
||||
class DummyChunk:
|
||||
def __init__(self, content, similarity):
|
||||
self.content = content
|
||||
self.similarity = similarity
|
||||
|
||||
|
||||
class DummyDocument:
|
||||
def __init__(self, id, title, chunks=None):
|
||||
self.id = id
|
||||
self.title = title
|
||||
self.chunks = chunks or []
|
||||
|
||||
|
||||
# Patch imports in dify.py to use dummy classes
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_imports(monkeypatch):
|
||||
import src.rag.dify as dify
|
||||
|
||||
dify.Resource = DummyResource
|
||||
dify.Chunk = DummyChunk
|
||||
dify.Document = DummyDocument
|
||||
yield
|
||||
|
||||
|
||||
def test_parse_uri_valid():
|
||||
uri = "rag://dataset/123#abc"
|
||||
dataset_id, document_id = parse_uri(uri)
|
||||
assert dataset_id == "123"
|
||||
assert document_id == "abc"
|
||||
|
||||
|
||||
def test_parse_uri_invalid():
|
||||
with pytest.raises(ValueError):
|
||||
parse_uri("http://dataset/123#abc")
|
||||
|
||||
|
||||
def test_init_env_vars(monkeypatch):
|
||||
monkeypatch.setenv("DIFY_API_URL", "http://api")
|
||||
monkeypatch.setenv("DIFY_API_KEY", "key")
|
||||
provider = DifyProvider()
|
||||
assert provider.api_url == "http://api"
|
||||
assert provider.api_key == "key"
|
||||
|
||||
|
||||
def test_init_missing_env(monkeypatch):
|
||||
monkeypatch.delenv("DIFY_API_URL", raising=False)
|
||||
monkeypatch.setenv("DIFY_API_KEY", "key")
|
||||
with pytest.raises(ValueError):
|
||||
DifyProvider()
|
||||
monkeypatch.setenv("DIFY_API_URL", "http://api")
|
||||
monkeypatch.delenv("DIFY_API_KEY", raising=False)
|
||||
with pytest.raises(ValueError):
|
||||
DifyProvider()
|
||||
|
||||
|
||||
@patch("src.rag.dify.requests.post")
|
||||
def test_query_relevant_documents_success(mock_post, monkeypatch):
|
||||
monkeypatch.setenv("DIFY_API_URL", "http://api")
|
||||
monkeypatch.setenv("DIFY_API_KEY", "key")
|
||||
provider = DifyProvider()
|
||||
resource = DummyResource("rag://dataset/123#doc456")
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"records": [
|
||||
{
|
||||
"segment": {
|
||||
"content": "chunk text",
|
||||
"document": {
|
||||
"id": "doc456",
|
||||
"name": "Doc Title",
|
||||
},
|
||||
},
|
||||
"score": 0.9,
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
docs = provider.query_relevant_documents("query", [resource])
|
||||
assert len(docs) == 1
|
||||
assert docs[0].id == "doc456"
|
||||
assert docs[0].title == "Doc Title"
|
||||
assert len(docs[0].chunks) == 1
|
||||
assert docs[0].chunks[0].content == "chunk text"
|
||||
assert docs[0].chunks[0].similarity == 0.9
|
||||
|
||||
|
||||
@patch("src.rag.dify.requests.post")
|
||||
def test_query_relevant_documents_error(mock_post, monkeypatch):
|
||||
monkeypatch.setenv("DIFY_API_URL", "http://api")
|
||||
monkeypatch.setenv("DIFY_API_KEY", "key")
|
||||
provider = DifyProvider()
|
||||
resource = DummyResource("rag://dataset/123#doc456")
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "error"
|
||||
mock_post.return_value = mock_response
|
||||
with pytest.raises(Exception):
|
||||
provider.query_relevant_documents("query", [resource])
|
||||
|
||||
|
||||
@patch("src.rag.dify.requests.get")
|
||||
def test_list_resources_success(mock_get, monkeypatch):
|
||||
monkeypatch.setenv("DIFY_API_URL", "http://api")
|
||||
monkeypatch.setenv("DIFY_API_KEY", "key")
|
||||
provider = DifyProvider()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"data": [
|
||||
{"id": "123", "name": "Dataset1", "description": "desc1"},
|
||||
{"id": "456", "name": "Dataset2", "description": "desc2"},
|
||||
]
|
||||
}
|
||||
mock_get.return_value = mock_response
|
||||
resources = provider.list_resources()
|
||||
assert len(resources) == 2
|
||||
assert resources[0].uri == "rag://dataset/123"
|
||||
assert resources[0].title == "Dataset1"
|
||||
assert resources[0].description == "desc1"
|
||||
assert resources[1].uri == "rag://dataset/456"
|
||||
assert resources[1].title == "Dataset2"
|
||||
assert resources[1].description == "desc2"
|
||||
|
||||
|
||||
@patch("src.rag.dify.requests.get")
|
||||
def test_list_resources_error(mock_get, monkeypatch):
|
||||
monkeypatch.setenv("DIFY_API_URL", "http://api")
|
||||
monkeypatch.setenv("DIFY_API_KEY", "key")
|
||||
provider = DifyProvider()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.text = "fail"
|
||||
mock_get.return_value = mock_response
|
||||
with pytest.raises(Exception):
|
||||
provider.list_resources()
|
||||
@@ -0,0 +1,930 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
"""
|
||||
Tests for Milvus RAG provider.
|
||||
|
||||
IMPORTANT NOTE: This test file creates temporary directories for testing examples
|
||||
functionality. All temporary directories are automatically cleaned up using pytest
|
||||
fixtures. When adding new tests that create temporary directories:
|
||||
|
||||
1. Use the provided fixtures (temp_examples_dir, temp_error_examples_dir, etc.)
|
||||
2. Never create temporary directories without automatic cleanup
|
||||
3. Follow the pattern: fixture -> use -> automatic cleanup
|
||||
4. If you need a new directory pattern, create a corresponding fixture
|
||||
|
||||
This ensures tests don't leave behind temporary files that clutter the workspace.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
import src.rag.milvus as milvus_mod
|
||||
from src.rag.milvus import MilvusProvider
|
||||
from src.rag.retriever import Resource
|
||||
|
||||
|
||||
class DummyEmbedding:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
def embed_query(self, text: str):
|
||||
return [0.1, 0.2, 0.3]
|
||||
|
||||
def embed_documents(self, texts):
|
||||
return [[0.1, 0.2, 0.3] for _ in texts]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_embeddings(monkeypatch):
|
||||
# Prevent network / external API usage during __init__
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_PROVIDER", "openai")
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_MODEL", "text-embedding-ada-002")
|
||||
monkeypatch.setenv("MILVUS_COLLECTION", "documents")
|
||||
monkeypatch.setenv("MILVUS_URI", "./milvus_demo.db") # default lite
|
||||
monkeypatch.setattr(milvus_mod, "OpenAIEmbeddings", DummyEmbedding)
|
||||
monkeypatch.setattr(milvus_mod, "DashscopeEmbeddings", DummyEmbedding)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def project_root():
|
||||
# Mirror logic from implementation: current_file.parent.parent.parent
|
||||
return Path(milvus_mod.__file__).parent.parent.parent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_examples_dir(project_root):
|
||||
"""Create a temporary examples directory with automatic cleanup."""
|
||||
# Create a unique temporary directory name
|
||||
temp_dir_name = f"examples_test_{uuid4().hex}"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
|
||||
# Create the directory
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
yield temp_dir_path
|
||||
|
||||
# Cleanup: remove the directory and all its contents
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_error_examples_dir(project_root):
|
||||
"""Create a temporary error examples directory with automatic cleanup."""
|
||||
# Create a unique temporary directory name for error tests
|
||||
temp_dir_name = f"examples_error_{uuid4().hex}"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
|
||||
# Create the directory
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
yield temp_dir_path
|
||||
|
||||
# Cleanup: remove the directory and all its contents
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_load_skip_examples_dir(project_root):
|
||||
"""Create a temporary examples directory for load_skip tests with automatic cleanup."""
|
||||
# Use the expected directory name for this test
|
||||
temp_dir_name = "examples_test_load_skip"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
|
||||
# Create the directory if it doesn't exist
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
yield temp_dir_path
|
||||
|
||||
# Cleanup: remove the directory and all its contents
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_single_chunk_examples_dir(project_root):
|
||||
"""Create a temporary examples directory for single_chunk tests with automatic cleanup."""
|
||||
# Use the expected directory name for this test
|
||||
temp_dir_name = "examples_test_single_chunk"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
|
||||
# Create the directory if it doesn't exist
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
yield temp_dir_path
|
||||
|
||||
# Cleanup: remove the directory and all its contents
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
def _patch_init(monkeypatch):
|
||||
"""Patch retriever initialization to use dummy embedding model."""
|
||||
monkeypatch.setattr(
|
||||
MilvusProvider,
|
||||
"_init_embedding_model",
|
||||
lambda self: setattr(self, "embedding_model", DummyEmbedding()),
|
||||
)
|
||||
|
||||
|
||||
def test_list_local_markdown_resources_missing_dir(project_root):
|
||||
retriever = MilvusProvider()
|
||||
# Point to a non-existent examples dir
|
||||
retriever.examples_dir = f"missing_examples_{uuid4().hex}"
|
||||
resources = retriever._list_local_markdown_resources()
|
||||
assert resources == []
|
||||
|
||||
|
||||
def test_list_local_markdown_resources_populated(temp_examples_dir):
|
||||
retriever = MilvusProvider()
|
||||
# Use the name of the temp directory for examples_dir
|
||||
retriever.examples_dir = temp_examples_dir.name
|
||||
|
||||
# File with heading
|
||||
(temp_examples_dir / "file1.md").write_text(
|
||||
"# Title One\n\nContent body.", encoding="utf-8"
|
||||
)
|
||||
# File without heading -> fallback title
|
||||
(temp_examples_dir / "file_two.md").write_text("No heading here.", encoding="utf-8")
|
||||
# Non-markdown file should be ignored
|
||||
(temp_examples_dir / "ignore.txt").write_text(
|
||||
"Should not be picked up.", encoding="utf-8"
|
||||
)
|
||||
|
||||
resources = retriever._list_local_markdown_resources()
|
||||
# Order not guaranteed; sort by uri for assertions
|
||||
resources.sort(key=lambda r: r.uri)
|
||||
|
||||
# Expect two resources
|
||||
assert len(resources) == 2
|
||||
uris = {r.uri for r in resources}
|
||||
assert uris == {
|
||||
f"milvus://{retriever.collection_name}/file1.md",
|
||||
f"milvus://{retriever.collection_name}/file_two.md",
|
||||
}
|
||||
|
||||
res_map = {r.uri: r for r in resources}
|
||||
r1 = res_map[f"milvus://{retriever.collection_name}/file1.md"]
|
||||
assert isinstance(r1, Resource)
|
||||
assert r1.title == "Title One"
|
||||
assert r1.description == "Local markdown example (not yet ingested)"
|
||||
|
||||
r2 = res_map[f"milvus://{retriever.collection_name}/file_two.md"]
|
||||
# Fallback logic: filename -> "file_two" -> "file two" -> title case -> "File Two"
|
||||
assert r2.title == "File Two"
|
||||
assert r2.description == "Local markdown example (not yet ingested)"
|
||||
|
||||
|
||||
def test_list_local_markdown_resources_read_error(monkeypatch, temp_error_examples_dir):
|
||||
retriever = MilvusProvider()
|
||||
# Use the name of the temp directory for examples_dir
|
||||
retriever.examples_dir = temp_error_examples_dir.name
|
||||
|
||||
bad_file = temp_error_examples_dir / "bad.md"
|
||||
good_file = temp_error_examples_dir / "good.md"
|
||||
good_file.write_text("# Good Title\n\nBody.", encoding="utf-8")
|
||||
bad_file.write_text("Broken", encoding="utf-8")
|
||||
|
||||
# Patch Path.read_text to raise for bad.md only
|
||||
original_read_text = Path.read_text
|
||||
|
||||
def fake_read_text(self, *args, **kwargs):
|
||||
if self == bad_file:
|
||||
raise OSError("Cannot read file")
|
||||
return original_read_text(self, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(Path, "read_text", fake_read_text)
|
||||
|
||||
resources = retriever._list_local_markdown_resources()
|
||||
# Only good.md should appear
|
||||
assert len(resources) == 1
|
||||
r = resources[0]
|
||||
assert r.title == "Good Title"
|
||||
assert r.uri == f"milvus://{retriever.collection_name}/good.md"
|
||||
|
||||
|
||||
def test_create_collection_schema_fields(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
schema = retriever._create_collection_schema()
|
||||
field_names = {f.name for f in schema.fields}
|
||||
# Core fields must be present
|
||||
assert {
|
||||
retriever.id_field,
|
||||
retriever.vector_field,
|
||||
retriever.content_field,
|
||||
} <= field_names
|
||||
# Dynamic field enabled for extra metadata
|
||||
assert schema.enable_dynamic_field is True
|
||||
|
||||
|
||||
def test_generate_doc_id_stable(monkeypatch, tmp_path):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
test_file = tmp_path / "example.md"
|
||||
test_file.write_text("# Title\nBody", encoding="utf-8")
|
||||
doc_id1 = retriever._generate_doc_id(test_file)
|
||||
doc_id2 = retriever._generate_doc_id(test_file)
|
||||
assert doc_id1 == doc_id2 # deterministic given unchanged file metadata
|
||||
|
||||
|
||||
def test_extract_title_from_markdown(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
heading = retriever._extract_title_from_markdown("# Heading\nBody", "ignored.md")
|
||||
assert heading == "Heading"
|
||||
fallback = retriever._extract_title_from_markdown("Body only", "my_file_name.md")
|
||||
assert fallback == "My File Name"
|
||||
|
||||
|
||||
def test_split_content_chunking(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_CHUNK_SIZE", "40") # small to force split
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
long_content = (
|
||||
"Para1 text here.\n\nPara2 second block.\n\nPara3 final." # 3 paragraphs
|
||||
)
|
||||
chunks = retriever._split_content(long_content)
|
||||
assert len(chunks) >= 2 # forced split
|
||||
assert all(chunks) # no empty chunks
|
||||
|
||||
|
||||
def test_get_embedding_invalid_inputs(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
# Non-string value
|
||||
with pytest.raises(RuntimeError):
|
||||
retriever._get_embedding(123) # type: ignore[arg-type]
|
||||
# Whitespace only
|
||||
with pytest.raises(RuntimeError):
|
||||
retriever._get_embedding(" ")
|
||||
|
||||
|
||||
def test_list_resources_remote_success_and_dedup(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
class DocObj:
|
||||
def __init__(self, content: str, meta: dict):
|
||||
self.page_content = content
|
||||
self.metadata = meta
|
||||
|
||||
calls = {"similarity_search": 0}
|
||||
|
||||
class RemoteClient:
|
||||
def similarity_search(self, query, k, expr): # noqa: D401
|
||||
calls["similarity_search"] += 1
|
||||
# Two docs with identical id to test dedup
|
||||
meta1 = {
|
||||
retriever.id_field: "d1",
|
||||
retriever.title_field: "T1",
|
||||
retriever.url_field: "u1",
|
||||
}
|
||||
meta2 = {
|
||||
retriever.id_field: "d1",
|
||||
retriever.title_field: "T1_dup",
|
||||
retriever.url_field: "u1",
|
||||
}
|
||||
return [DocObj("c1", meta1), DocObj("c1_dup", meta2)]
|
||||
|
||||
retriever.client = RemoteClient()
|
||||
resources = retriever.list_resources("query text")
|
||||
assert len(resources) == 1 # dedup applied
|
||||
assert resources[0].title.startswith("T1")
|
||||
assert calls["similarity_search"] == 1
|
||||
|
||||
|
||||
def test_list_resources_lite_success(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
class DummyMilvusLite:
|
||||
def query(self, collection_name, filter, output_fields, limit): # noqa: D401
|
||||
return [
|
||||
{
|
||||
retriever.id_field: "idA",
|
||||
retriever.title_field: "Alpha",
|
||||
retriever.url_field: "u://a",
|
||||
},
|
||||
{
|
||||
retriever.id_field: "idB",
|
||||
retriever.title_field: "Beta",
|
||||
retriever.url_field: "u://b",
|
||||
},
|
||||
]
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
resources = retriever.list_resources()
|
||||
assert {r.title for r in resources} == {"Alpha", "Beta"}
|
||||
|
||||
|
||||
def test_query_relevant_documents_lite_success(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
# Provide deterministic embedding output
|
||||
retriever.embedding_model.embed_query = lambda text: [0.1, 0.2, 0.3] # type: ignore
|
||||
|
||||
class DummyMilvusLite:
|
||||
def search(
|
||||
self, collection_name, data, anns_field, param, limit, output_fields
|
||||
): # noqa: D401
|
||||
# Simulate two result entries
|
||||
return [
|
||||
[
|
||||
{
|
||||
"entity": {
|
||||
retriever.id_field: "d1",
|
||||
retriever.content_field: "c1",
|
||||
retriever.title_field: "T1",
|
||||
retriever.url_field: "u1",
|
||||
},
|
||||
"distance": 0.9,
|
||||
},
|
||||
{
|
||||
"entity": {
|
||||
retriever.id_field: "d2",
|
||||
retriever.content_field: "c2",
|
||||
retriever.title_field: "T2",
|
||||
retriever.url_field: "u2",
|
||||
},
|
||||
"distance": 0.8,
|
||||
},
|
||||
]
|
||||
]
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
# Filter for only d2 via resource list
|
||||
docs = retriever.query_relevant_documents(
|
||||
"question", resources=[Resource(uri="milvus://d2", title="", description="")]
|
||||
)
|
||||
assert len(docs) == 1 and docs[0].id == "d2" and docs[0].chunks[0].similarity == 0.8
|
||||
|
||||
|
||||
def test_query_relevant_documents_remote_success(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever.embedding_model.embed_query = lambda text: [0.1, 0.2, 0.3] # type: ignore
|
||||
|
||||
class DocObj:
|
||||
def __init__(self, content: str, meta: dict): # noqa: D401
|
||||
self.page_content = content
|
||||
self.metadata = meta
|
||||
|
||||
class RemoteClient:
|
||||
def similarity_search_with_score(self, query, k): # noqa: D401
|
||||
return [
|
||||
(
|
||||
DocObj(
|
||||
"c1",
|
||||
{
|
||||
retriever.id_field: "d1",
|
||||
retriever.title_field: "T1",
|
||||
retriever.url_field: "u1",
|
||||
},
|
||||
),
|
||||
0.7,
|
||||
),
|
||||
(
|
||||
DocObj(
|
||||
"c2",
|
||||
{
|
||||
retriever.id_field: "d2",
|
||||
retriever.title_field: "T2",
|
||||
retriever.url_field: "u2",
|
||||
},
|
||||
),
|
||||
0.6,
|
||||
),
|
||||
]
|
||||
|
||||
retriever.client = RemoteClient()
|
||||
# Filter to only d1
|
||||
docs = retriever.query_relevant_documents(
|
||||
"q", resources=[Resource(uri="milvus://d1", title="", description="")]
|
||||
)
|
||||
assert len(docs) == 1 and docs[0].id == "d1" and docs[0].chunks[0].similarity == 0.7
|
||||
|
||||
|
||||
def test_get_embedding_dimension_explicit(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_DIM", "777")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
assert retriever.embedding_dim == 777
|
||||
|
||||
|
||||
def test_get_embedding_dimension_unknown_model(monkeypatch):
|
||||
monkeypatch.delenv("MILVUS_EMBEDDING_DIM", raising=False)
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_MODEL", "unknown-model-x")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
# falls back to default 1536
|
||||
assert retriever.embedding_dim == 1536
|
||||
|
||||
|
||||
def test_is_milvus_lite_variants(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
monkeypatch.setenv("MILVUS_URI", "mydb.db")
|
||||
assert MilvusProvider()._is_milvus_lite() is True
|
||||
monkeypatch.setenv("MILVUS_URI", "relative_path_store")
|
||||
assert MilvusProvider()._is_milvus_lite() is True
|
||||
monkeypatch.setenv("MILVUS_URI", "http://host:19530")
|
||||
assert MilvusProvider()._is_milvus_lite() is False
|
||||
|
||||
|
||||
def test_create_collection_lite(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
created: dict = {}
|
||||
|
||||
class DummyMilvusLite:
|
||||
def list_collections(self): # noqa: D401
|
||||
return [] # empty triggers creation
|
||||
|
||||
def create_collection(self, collection_name, schema, index_params): # noqa: D401
|
||||
created["name"] = collection_name
|
||||
created["schema"] = schema
|
||||
created["index"] = index_params
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
retriever._ensure_collection_exists()
|
||||
assert created["name"] == retriever.collection_name
|
||||
|
||||
|
||||
def test_ensure_collection_exists_remote(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote:19530")
|
||||
retriever = MilvusProvider()
|
||||
# remote path, nothing thrown
|
||||
retriever.client = SimpleNamespace()
|
||||
retriever._ensure_collection_exists()
|
||||
|
||||
|
||||
def test_get_existing_document_ids_lite(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
class DummyMilvusLite:
|
||||
def query(self, collection_name, filter, output_fields, limit): # noqa: D401
|
||||
return [
|
||||
{retriever.id_field: "a"},
|
||||
{retriever.id_field: "b"},
|
||||
{"other": "ignored"},
|
||||
]
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
assert retriever._get_existing_document_ids() == {"a", "b"}
|
||||
|
||||
|
||||
def test_get_existing_document_ids_remote(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
monkeypatch.setenv("MILVUS_URI", "http://x")
|
||||
retriever = MilvusProvider()
|
||||
retriever.client = object()
|
||||
assert retriever._get_existing_document_ids() == set()
|
||||
|
||||
|
||||
def test_insert_document_chunk_lite_and_error(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
captured = {}
|
||||
|
||||
class DummyMilvusLite:
|
||||
def insert(self, collection_name, data): # noqa: D401
|
||||
captured["data"] = data
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
retriever._insert_document_chunk(
|
||||
doc_id="id1", content="hello", title="T", url="u", metadata={"m": 1}
|
||||
)
|
||||
assert captured["data"][0][retriever.id_field] == "id1"
|
||||
|
||||
# error path: patch embedding to raise
|
||||
def bad_embed(text): # noqa: D401
|
||||
raise RuntimeError("boom")
|
||||
|
||||
retriever.embedding_model.embed_query = bad_embed # type: ignore[attr-defined]
|
||||
with pytest.raises(RuntimeError):
|
||||
retriever._insert_document_chunk(
|
||||
doc_id="id2", content="err", title="T", url="u", metadata={}
|
||||
)
|
||||
|
||||
|
||||
def test_insert_document_chunk_remote(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
retriever = MilvusProvider()
|
||||
added = {}
|
||||
|
||||
class RemoteClient:
|
||||
def add_texts(self, texts, metadatas): # noqa: D401
|
||||
added["texts"] = texts
|
||||
added["meta"] = metadatas
|
||||
|
||||
retriever.client = RemoteClient()
|
||||
retriever._insert_document_chunk(
|
||||
doc_id="idx", content="ct", title="Title", url="urlx", metadata={"k": 2}
|
||||
)
|
||||
assert added["meta"][0][retriever.id_field] == "idx"
|
||||
|
||||
|
||||
def test_connect_lite_and_error(monkeypatch):
|
||||
# patch MilvusClient to a dummy
|
||||
class FakeMilvusClient:
|
||||
def __init__(self, uri): # noqa: D401
|
||||
self.uri = uri
|
||||
|
||||
def list_collections(self): # noqa: D401
|
||||
return []
|
||||
|
||||
def create_collection(self, **kwargs): # noqa: D401
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(milvus_mod, "MilvusClient", FakeMilvusClient)
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever._connect()
|
||||
assert isinstance(retriever.client, FakeMilvusClient)
|
||||
|
||||
# error path: patch MilvusClient to raise
|
||||
class BadClient:
|
||||
def __init__(self, uri): # noqa: D401
|
||||
raise RuntimeError("fail connect")
|
||||
|
||||
monkeypatch.setattr(milvus_mod, "MilvusClient", BadClient)
|
||||
retriever2 = MilvusProvider()
|
||||
with pytest.raises(ConnectionError):
|
||||
retriever2._connect()
|
||||
|
||||
|
||||
def test_connect_remote(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
_patch_init(monkeypatch)
|
||||
created = {}
|
||||
|
||||
class FakeLangchainMilvus:
|
||||
def __init__(self, **kwargs): # noqa: D401
|
||||
created.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(milvus_mod, "LangchainMilvus", FakeLangchainMilvus)
|
||||
retriever = MilvusProvider()
|
||||
retriever._connect()
|
||||
assert created["collection_name"] == retriever.collection_name
|
||||
|
||||
|
||||
def test_list_resources_remote_failure(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
# Provide minimal working local examples dir (none -> returns [])
|
||||
monkeypatch.setattr(retriever, "_list_local_markdown_resources", lambda: [])
|
||||
|
||||
# patch client to raise inside similarity_search to trigger fallback path
|
||||
class BadClient:
|
||||
def similarity_search(self, *args, **kwargs): # noqa: D401
|
||||
raise RuntimeError("fail")
|
||||
|
||||
retriever.client = BadClient()
|
||||
# Should fallback to [] without raising
|
||||
assert retriever.list_resources() == []
|
||||
|
||||
|
||||
def test_list_local_markdown_resources_empty(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
monkeypatch.setenv("MILVUS_EXAMPLES_DIR", "nonexistent_dir")
|
||||
retriever.examples_dir = "nonexistent_dir"
|
||||
assert retriever._list_local_markdown_resources() == []
|
||||
|
||||
|
||||
def test_query_relevant_documents_error(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever.embedding_model.embed_query = lambda text: ( # type: ignore
|
||||
_ for _ in ()
|
||||
).throw(RuntimeError("embed fail"))
|
||||
with pytest.raises(RuntimeError):
|
||||
retriever.query_relevant_documents("q")
|
||||
|
||||
|
||||
def test_create_collection_when_client_exists(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever.client = SimpleNamespace(closed=False)
|
||||
# remote vs lite path difference handled by _is_milvus_lite
|
||||
retriever.create_collection() # should no-op gracefully
|
||||
|
||||
|
||||
def test_load_examples_force_reload(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever.client = SimpleNamespace()
|
||||
called = {"clear": 0, "load": 0}
|
||||
monkeypatch.setattr(
|
||||
retriever, "_clear_example_documents", lambda: called.__setitem__("clear", 1)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
retriever, "_load_example_files", lambda: called.__setitem__("load", 1)
|
||||
)
|
||||
retriever.load_examples(force_reload=True)
|
||||
assert called == {"clear": 1, "load": 1}
|
||||
|
||||
|
||||
def test_clear_example_documents_remote(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever.client = SimpleNamespace()
|
||||
# Should just log and not raise
|
||||
retriever._clear_example_documents()
|
||||
|
||||
|
||||
def test_clear_example_documents_lite(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
deleted = {}
|
||||
|
||||
class DummyMilvusLite:
|
||||
def query(self, **kwargs): # noqa: D401
|
||||
return [
|
||||
{retriever.id_field: "ex1"},
|
||||
{retriever.id_field: "ex2"},
|
||||
]
|
||||
|
||||
def delete(self, collection_name, ids): # noqa: D401
|
||||
deleted["ids"] = ids
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
retriever._clear_example_documents()
|
||||
assert deleted["ids"] == ["ex1", "ex2"]
|
||||
|
||||
|
||||
def test_get_loaded_examples_lite_and_error(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
|
||||
class DummyMilvusLite:
|
||||
def query(self, **kwargs): # noqa: D401
|
||||
return [
|
||||
{
|
||||
retriever.id_field: "id1",
|
||||
retriever.title_field: "T1",
|
||||
retriever.url_field: "u1",
|
||||
"file": "f1",
|
||||
}
|
||||
]
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
loaded = retriever.get_loaded_examples()
|
||||
assert loaded[0]["id"] == "id1"
|
||||
|
||||
# error path
|
||||
class BadClient:
|
||||
def query(self, **kwargs): # noqa: D401
|
||||
raise RuntimeError("fail")
|
||||
|
||||
retriever.client = BadClient()
|
||||
assert retriever.get_loaded_examples() == []
|
||||
|
||||
|
||||
def test_get_loaded_examples_remote(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
retriever.client = SimpleNamespace()
|
||||
assert retriever.get_loaded_examples() == []
|
||||
|
||||
|
||||
def test_close_lite_and_remote(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
closed = {"c": 0}
|
||||
|
||||
class DummyMilvusLite:
|
||||
def close(self): # noqa: D401
|
||||
closed["c"] += 1
|
||||
|
||||
def list_collections(self): # noqa: D401
|
||||
return []
|
||||
|
||||
def create_collection(self, **kwargs): # noqa: D401
|
||||
pass
|
||||
|
||||
retriever.client = DummyMilvusLite()
|
||||
retriever.close()
|
||||
assert closed["c"] == 1
|
||||
|
||||
# remote path: no close attr usage expected
|
||||
monkeypatch.setenv("MILVUS_URI", "http://remote")
|
||||
retriever2 = MilvusProvider()
|
||||
retriever2.client = SimpleNamespace()
|
||||
retriever2.close() # should not raise
|
||||
|
||||
|
||||
def test_get_embedding_invalid_output(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
retriever = MilvusProvider()
|
||||
# patch embedding model to return invalid output (empty list)
|
||||
retriever.embedding_model.embed_query = lambda text: [] # type: ignore
|
||||
with pytest.raises(RuntimeError):
|
||||
retriever._get_embedding("text")
|
||||
|
||||
|
||||
def test_dashscope_embeddings_empty_inputs_short_circuit(monkeypatch):
|
||||
# Use real class but swap _client to ensure create is never called
|
||||
emb = milvus_mod.DashscopeEmbeddings(model="m")
|
||||
|
||||
class FailingClient:
|
||||
class _Emb:
|
||||
def create(self, *a, **k):
|
||||
raise AssertionError("Should not be called for empty input")
|
||||
|
||||
embeddings = _Emb()
|
||||
|
||||
emb._client = FailingClient() # type: ignore
|
||||
assert emb.embed_documents([]) == []
|
||||
|
||||
|
||||
# Tests for _init_embedding_model provider selection logic
|
||||
def test_init_embedding_model_openai(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_PROVIDER", "openai")
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_MODEL", "text-embedding-ada-002")
|
||||
captured = {}
|
||||
|
||||
class CapturingOpenAI:
|
||||
def __init__(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(milvus_mod, "OpenAIEmbeddings", CapturingOpenAI)
|
||||
prov = MilvusProvider()
|
||||
assert isinstance(prov.embedding_model, CapturingOpenAI)
|
||||
# kwargs forwarded
|
||||
assert captured["model"] == "text-embedding-ada-002"
|
||||
assert captured["encoding_format"] == "float"
|
||||
assert captured["dimensions"] == prov.embedding_dim
|
||||
|
||||
|
||||
def test_init_embedding_model_dashscope(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_PROVIDER", "dashscope")
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_MODEL", "text-embedding-ada-002")
|
||||
captured = {}
|
||||
|
||||
class CapturingDashscope:
|
||||
def __init__(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
|
||||
monkeypatch.setattr(milvus_mod, "DashscopeEmbeddings", CapturingDashscope)
|
||||
prov = MilvusProvider()
|
||||
assert isinstance(prov.embedding_model, CapturingDashscope)
|
||||
assert captured["model"] == "text-embedding-ada-002"
|
||||
assert captured["encoding_format"] == "float"
|
||||
assert captured["dimensions"] == prov.embedding_dim
|
||||
|
||||
|
||||
def test_init_embedding_model_invalid_provider(monkeypatch):
|
||||
monkeypatch.setenv("MILVUS_EMBEDDING_PROVIDER", "not_a_provider")
|
||||
with pytest.raises(ValueError):
|
||||
MilvusProvider()
|
||||
|
||||
|
||||
def test_load_example_files_directory_missing(monkeypatch):
|
||||
_patch_init(monkeypatch)
|
||||
missing_dir = "examples_dir_does_not_exist_xyz"
|
||||
monkeypatch.setenv("MILVUS_EXAMPLES_DIR", missing_dir)
|
||||
retriever = MilvusProvider()
|
||||
retriever.examples_dir = missing_dir
|
||||
called = {"insert": 0}
|
||||
monkeypatch.setattr(
|
||||
retriever,
|
||||
"_insert_document_chunk",
|
||||
lambda **kwargs: (_ for _ in ()).throw(AssertionError("should not insert")),
|
||||
)
|
||||
retriever._load_example_files()
|
||||
assert called["insert"] == 0 # sanity (no insertion attempted)
|
||||
|
||||
|
||||
def test_load_example_files_loads_and_skips_existing(
|
||||
monkeypatch, temp_load_skip_examples_dir
|
||||
):
|
||||
_patch_init(monkeypatch)
|
||||
examples_dir_name = temp_load_skip_examples_dir.name
|
||||
|
||||
file1 = temp_load_skip_examples_dir / "file1.md"
|
||||
file2 = temp_load_skip_examples_dir / "file2.md"
|
||||
file1.write_text("# Title One\nContent A", encoding="utf-8")
|
||||
file2.write_text("# Title Two\nContent B", encoding="utf-8")
|
||||
|
||||
monkeypatch.setenv("MILVUS_EXAMPLES_DIR", examples_dir_name)
|
||||
retriever = MilvusProvider()
|
||||
retriever.examples_dir = examples_dir_name
|
||||
|
||||
# Compute doc ids using real method
|
||||
doc_id_file1 = retriever._generate_doc_id(file1)
|
||||
doc_id_file2 = retriever._generate_doc_id(file2)
|
||||
|
||||
# Existing docs contains file1 so it is skipped
|
||||
monkeypatch.setattr(retriever, "_get_existing_document_ids", lambda: {doc_id_file1})
|
||||
# Force two chunks for any file to test suffix logic
|
||||
monkeypatch.setattr(retriever, "_split_content", lambda content: ["part1", "part2"])
|
||||
|
||||
calls = []
|
||||
|
||||
def record_insert(doc_id, content, title, url, metadata):
|
||||
calls.append(
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"content": content,
|
||||
"title": title,
|
||||
"url": url,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(retriever, "_insert_document_chunk", record_insert)
|
||||
|
||||
retriever._load_example_files()
|
||||
|
||||
# Only file2 processed -> two chunk inserts
|
||||
assert len(calls) == 2
|
||||
expected_ids = {f"{doc_id_file2}_chunk_0", f"{doc_id_file2}_chunk_1"}
|
||||
assert {c["doc_id"] for c in calls} == expected_ids
|
||||
assert all(c["metadata"]["file"] == "file2.md" for c in calls)
|
||||
assert all(c["metadata"]["source"] == "examples" for c in calls)
|
||||
assert all(c["title"] == "Title Two" for c in calls)
|
||||
|
||||
|
||||
def test_load_example_files_single_chunk_no_suffix(
|
||||
monkeypatch, temp_single_chunk_examples_dir
|
||||
):
|
||||
_patch_init(monkeypatch)
|
||||
examples_dir_name = temp_single_chunk_examples_dir.name
|
||||
|
||||
file_single = temp_single_chunk_examples_dir / "single.md"
|
||||
file_single.write_text(
|
||||
"# Single Title\nOnly one small paragraph.", encoding="utf-8"
|
||||
)
|
||||
|
||||
monkeypatch.setenv("MILVUS_EXAMPLES_DIR", examples_dir_name)
|
||||
retriever = MilvusProvider()
|
||||
retriever.examples_dir = examples_dir_name
|
||||
|
||||
base_doc_id = retriever._generate_doc_id(file_single)
|
||||
|
||||
monkeypatch.setattr(retriever, "_get_existing_document_ids", lambda: set())
|
||||
monkeypatch.setattr(retriever, "_split_content", lambda content: ["onlychunk"])
|
||||
|
||||
captured = {}
|
||||
|
||||
def capture(doc_id, content, title, url, metadata):
|
||||
captured["doc_id"] = doc_id
|
||||
captured["title"] = title
|
||||
captured["metadata"] = metadata
|
||||
|
||||
monkeypatch.setattr(retriever, "_insert_document_chunk", capture)
|
||||
|
||||
retriever._load_example_files()
|
||||
|
||||
assert captured["doc_id"] == base_doc_id # no _chunk_ suffix
|
||||
assert captured["title"] == "Single Title"
|
||||
assert captured["metadata"]["file"] == "single.md"
|
||||
assert captured["metadata"]["source"] == "examples"
|
||||
|
||||
|
||||
# Clean up test database file after tests
|
||||
import atexit
|
||||
|
||||
|
||||
def cleanup_test_database():
|
||||
"""Clean up milvus_demo.db file created during testing."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Skip cleanup if disabled
|
||||
if os.getenv("DISABLE_TEST_CLEANUP", "false").lower() == "true":
|
||||
return
|
||||
|
||||
db_file = Path.cwd() / "milvus_demo.db"
|
||||
if db_file.exists():
|
||||
try:
|
||||
db_file.unlink()
|
||||
print("🧹 Cleaned up milvus_demo.db")
|
||||
except Exception:
|
||||
pass # Silently ignore cleanup errors
|
||||
|
||||
|
||||
# Register cleanup to run when Python exits
|
||||
atexit.register(cleanup_test_database)
|
||||
@@ -0,0 +1,333 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
import src.rag.qdrant as qdrant_mod
|
||||
from src.rag.qdrant import QdrantProvider
|
||||
|
||||
|
||||
class DummyEmbedding:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
def embed_query(self, text: str):
|
||||
return [0.1] * 1536
|
||||
|
||||
def embed_documents(self, texts):
|
||||
return [[0.1] * 1536 for _ in texts]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_embeddings(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_PROVIDER", "openai")
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_MODEL", "text-embedding-ada-002")
|
||||
monkeypatch.setenv("QDRANT_COLLECTION", "documents")
|
||||
monkeypatch.setenv("QDRANT_LOCATION", ":memory:")
|
||||
monkeypatch.setattr(qdrant_mod, "OpenAIEmbeddings", DummyEmbedding)
|
||||
monkeypatch.setattr(qdrant_mod, "DashscopeEmbeddings", DummyEmbedding)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def project_root():
|
||||
return Path(qdrant_mod.__file__).parent.parent.parent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_examples_dir(project_root):
|
||||
temp_dir_name = f"examples_test_{uuid4().hex}"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
yield temp_dir_path
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_error_examples_dir(project_root):
|
||||
temp_dir_name = f"examples_error_{uuid4().hex}"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
yield temp_dir_path
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_load_skip_examples_dir(project_root):
|
||||
temp_dir_name = f"examples_load_skip_{uuid4().hex}"
|
||||
temp_dir_path = project_root / temp_dir_name
|
||||
temp_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
yield temp_dir_path
|
||||
if temp_dir_path.exists():
|
||||
shutil.rmtree(temp_dir_path)
|
||||
|
||||
|
||||
def test_init_openai_provider(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_PROVIDER", "openai")
|
||||
provider = QdrantProvider()
|
||||
assert provider.embedding_provider == "openai"
|
||||
assert isinstance(provider.embedding_model, DummyEmbedding)
|
||||
|
||||
|
||||
def test_init_dashscope_provider(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_PROVIDER", "dashscope")
|
||||
provider = QdrantProvider()
|
||||
assert provider.embedding_provider == "dashscope"
|
||||
assert isinstance(provider.embedding_model, DummyEmbedding)
|
||||
|
||||
|
||||
def test_init_invalid_provider(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_PROVIDER", "invalid_provider")
|
||||
with pytest.raises(ValueError, match="Unsupported embedding provider"):
|
||||
QdrantProvider()
|
||||
|
||||
|
||||
def test_get_embedding_dimension_explicit(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_DIM", "2048")
|
||||
provider = QdrantProvider()
|
||||
assert provider.embedding_dim == 2048
|
||||
|
||||
|
||||
def test_get_embedding_dimension_default(monkeypatch):
|
||||
monkeypatch.delenv("QDRANT_EMBEDDING_DIM", raising=False)
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_MODEL", "text-embedding-ada-002")
|
||||
provider = QdrantProvider()
|
||||
assert provider.embedding_dim == 1536
|
||||
|
||||
|
||||
def test_get_embedding_dimension_unknown_model(monkeypatch):
|
||||
monkeypatch.delenv("QDRANT_EMBEDDING_DIM", raising=False)
|
||||
monkeypatch.setenv("QDRANT_EMBEDDING_MODEL", "unknown-model")
|
||||
provider = QdrantProvider()
|
||||
assert provider.embedding_dim == 1536
|
||||
|
||||
|
||||
def test_connect_memory_mode(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_LOCATION", ":memory:")
|
||||
provider = QdrantProvider()
|
||||
provider._connect()
|
||||
assert provider.client is not None
|
||||
|
||||
|
||||
def test_create_collection(monkeypatch):
|
||||
provider = QdrantProvider()
|
||||
provider.create_collection()
|
||||
assert provider.client is not None
|
||||
|
||||
|
||||
def test_extract_title_from_markdown():
|
||||
provider = QdrantProvider()
|
||||
content = "# Test Title\n\nSome content"
|
||||
title = provider._extract_title_from_markdown(content, "test.md")
|
||||
assert title == "Test Title"
|
||||
|
||||
|
||||
def test_extract_title_fallback():
|
||||
provider = QdrantProvider()
|
||||
content = "No title here"
|
||||
title = provider._extract_title_from_markdown(content, "test_file.md")
|
||||
assert title == "Test File"
|
||||
|
||||
|
||||
def test_split_content_short():
|
||||
provider = QdrantProvider()
|
||||
content = "Short content"
|
||||
chunks = provider._split_content(content)
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0] == content
|
||||
|
||||
|
||||
def test_split_content_long(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_CHUNK_SIZE", "20")
|
||||
provider = QdrantProvider()
|
||||
content = "Paragraph one here\n\nParagraph two here\n\nParagraph three here\n\nParagraph four here"
|
||||
chunks = provider._split_content(content)
|
||||
assert len(chunks) > 1
|
||||
|
||||
|
||||
def test_string_to_uuid():
|
||||
provider = QdrantProvider()
|
||||
uuid1 = provider._string_to_uuid("test")
|
||||
uuid2 = provider._string_to_uuid("test")
|
||||
assert uuid1 == uuid2
|
||||
|
||||
|
||||
def test_get_embedding():
|
||||
provider = QdrantProvider()
|
||||
embedding = provider._get_embedding("test text")
|
||||
assert len(embedding) == 1536
|
||||
assert all(isinstance(x, float) for x in embedding)
|
||||
|
||||
|
||||
def test_load_examples_no_directory(monkeypatch, project_root):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", "nonexistent_dir")
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
|
||||
def test_load_examples_empty_directory(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
|
||||
def test_load_examples_with_files(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
|
||||
md_file = temp_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
loaded = provider.get_loaded_examples()
|
||||
assert len(loaded) == 1
|
||||
assert loaded[0]["title"] == "Test"
|
||||
|
||||
|
||||
def test_load_examples_skip_existing(monkeypatch, temp_load_skip_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_load_skip_examples_dir.name)
|
||||
|
||||
md_file = temp_load_skip_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
provider.load_examples()
|
||||
|
||||
loaded = provider.get_loaded_examples()
|
||||
assert len(loaded) == 1
|
||||
|
||||
|
||||
def test_load_examples_force_reload(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
|
||||
md_file = temp_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
provider.load_examples(force_reload=True)
|
||||
|
||||
loaded = provider.get_loaded_examples()
|
||||
assert len(loaded) == 1
|
||||
|
||||
|
||||
def test_load_examples_error_handling(monkeypatch, temp_error_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_error_examples_dir.name)
|
||||
|
||||
good_file = temp_error_examples_dir / "good.md"
|
||||
good_file.write_text("# Good\n\nContent", encoding="utf-8")
|
||||
|
||||
bad_file = temp_error_examples_dir / "bad.md"
|
||||
bad_file.write_text("# Bad\n\n", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
loaded = provider.get_loaded_examples()
|
||||
assert len(loaded) >= 1
|
||||
|
||||
|
||||
def test_list_resources_no_query(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
|
||||
md_file = temp_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
resources = provider.list_resources()
|
||||
assert len(resources) >= 1
|
||||
|
||||
|
||||
def test_list_resources_with_query(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
|
||||
md_file = temp_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
resources = provider.list_resources(query="test")
|
||||
assert isinstance(resources, list)
|
||||
|
||||
|
||||
def test_query_relevant_documents(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
|
||||
md_file = temp_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent about testing", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
documents = provider.query_relevant_documents("testing")
|
||||
assert isinstance(documents, list)
|
||||
|
||||
|
||||
def test_query_relevant_documents_with_resources(monkeypatch, temp_examples_dir):
|
||||
monkeypatch.setenv("QDRANT_EXAMPLES_DIR", temp_examples_dir.name)
|
||||
|
||||
md_file = temp_examples_dir / "test.md"
|
||||
md_file.write_text("# Test\n\nContent", encoding="utf-8")
|
||||
|
||||
provider = QdrantProvider()
|
||||
provider.load_examples()
|
||||
|
||||
resources = provider.list_resources()
|
||||
documents = provider.query_relevant_documents("test", resources=resources)
|
||||
assert isinstance(documents, list)
|
||||
|
||||
|
||||
def test_close():
|
||||
provider = QdrantProvider()
|
||||
provider._connect()
|
||||
provider.close()
|
||||
assert provider.client is None
|
||||
|
||||
|
||||
def test_del():
|
||||
provider = QdrantProvider()
|
||||
provider._connect()
|
||||
del provider
|
||||
|
||||
|
||||
def test_top_k_configuration(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_TOP_K", "20")
|
||||
provider = QdrantProvider()
|
||||
assert provider.top_k == 20
|
||||
|
||||
|
||||
def test_top_k_invalid(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_TOP_K", "invalid")
|
||||
provider = QdrantProvider()
|
||||
assert provider.top_k == 10
|
||||
|
||||
|
||||
def test_chunk_size_configuration(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_CHUNK_SIZE", "5000")
|
||||
provider = QdrantProvider()
|
||||
assert provider.chunk_size == 5000
|
||||
|
||||
|
||||
def test_collection_name_configuration(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_COLLECTION", "custom_collection")
|
||||
provider = QdrantProvider()
|
||||
assert provider.collection_name == "custom_collection"
|
||||
|
||||
|
||||
def test_auto_load_examples_configuration(monkeypatch):
|
||||
monkeypatch.setenv("QDRANT_AUTO_LOAD_EXAMPLES", "false")
|
||||
provider = QdrantProvider()
|
||||
assert provider.auto_load_examples is False
|
||||
@@ -0,0 +1,165 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.rag.ragflow import RAGFlowProvider, parse_uri
|
||||
|
||||
|
||||
# Dummy classes to mock dependencies
|
||||
class DummyResource:
|
||||
def __init__(self, uri, title="", description=""):
|
||||
self.uri = uri
|
||||
self.title = title
|
||||
self.description = description
|
||||
|
||||
|
||||
class DummyChunk:
|
||||
def __init__(self, content, similarity):
|
||||
self.content = content
|
||||
self.similarity = similarity
|
||||
|
||||
|
||||
class DummyDocument:
|
||||
def __init__(self, id, title, chunks=None):
|
||||
self.id = id
|
||||
self.title = title
|
||||
self.chunks = chunks or []
|
||||
|
||||
|
||||
# Patch imports in ragflow.py to use dummy classes
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_imports(monkeypatch):
|
||||
import src.rag.ragflow as ragflow
|
||||
|
||||
ragflow.Resource = DummyResource
|
||||
ragflow.Chunk = DummyChunk
|
||||
ragflow.Document = DummyDocument
|
||||
yield
|
||||
|
||||
|
||||
def test_parse_uri_valid():
|
||||
uri = "rag://dataset/123#abc"
|
||||
dataset_id, document_id = parse_uri(uri)
|
||||
assert dataset_id == "123"
|
||||
assert document_id == "abc"
|
||||
|
||||
|
||||
def test_parse_uri_invalid():
|
||||
with pytest.raises(ValueError):
|
||||
parse_uri("http://dataset/123#abc")
|
||||
|
||||
|
||||
def test_init_env_vars(monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
monkeypatch.delenv("RAGFLOW_PAGE_SIZE", raising=False)
|
||||
provider = RAGFlowProvider()
|
||||
assert provider.api_url == "http://api"
|
||||
assert provider.api_key == "key"
|
||||
assert provider.page_size == 10
|
||||
|
||||
|
||||
def test_init_page_size(monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
monkeypatch.setenv("RAGFLOW_PAGE_SIZE", "5")
|
||||
provider = RAGFlowProvider()
|
||||
assert provider.page_size == 5
|
||||
|
||||
|
||||
def test_init_cross_language(monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
monkeypatch.setenv("RAGFLOW_CROSS_LANGUAGES", "lang1,lang2")
|
||||
provider = RAGFlowProvider()
|
||||
assert provider.cross_languages == ["lang1", "lang2"]
|
||||
|
||||
|
||||
def test_init_missing_env(monkeypatch):
|
||||
monkeypatch.delenv("RAGFLOW_API_URL", raising=False)
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
with pytest.raises(ValueError):
|
||||
RAGFlowProvider()
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.delenv("RAGFLOW_API_KEY", raising=False)
|
||||
with pytest.raises(ValueError):
|
||||
RAGFlowProvider()
|
||||
|
||||
|
||||
@patch("src.rag.ragflow.requests.post")
|
||||
def test_query_relevant_documents_success(mock_post, monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
provider = RAGFlowProvider()
|
||||
resource = DummyResource("rag://dataset/123#doc456")
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"data": {
|
||||
"doc_aggs": [{"doc_id": "doc456", "doc_name": "Doc Title"}],
|
||||
"chunks": [
|
||||
{"document_id": "doc456", "content": "chunk text", "similarity": 0.9}
|
||||
],
|
||||
}
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
docs = provider.query_relevant_documents("query", [resource])
|
||||
assert len(docs) == 1
|
||||
assert docs[0].id == "doc456"
|
||||
assert docs[0].title == "Doc Title"
|
||||
assert len(docs[0].chunks) == 1
|
||||
assert docs[0].chunks[0].content == "chunk text"
|
||||
assert docs[0].chunks[0].similarity == 0.9
|
||||
|
||||
|
||||
@patch("src.rag.ragflow.requests.post")
|
||||
def test_query_relevant_documents_error(mock_post, monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
provider = RAGFlowProvider()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
mock_response.text = "error"
|
||||
mock_post.return_value = mock_response
|
||||
with pytest.raises(Exception):
|
||||
provider.query_relevant_documents("query", [])
|
||||
|
||||
|
||||
@patch("src.rag.ragflow.requests.get")
|
||||
def test_list_resources_success(mock_get, monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
provider = RAGFlowProvider()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"data": [
|
||||
{"id": "123", "name": "Dataset1", "description": "desc1"},
|
||||
{"id": "456", "name": "Dataset2", "description": "desc2"},
|
||||
]
|
||||
}
|
||||
mock_get.return_value = mock_response
|
||||
resources = provider.list_resources()
|
||||
assert len(resources) == 2
|
||||
assert resources[0].uri == "rag://dataset/123"
|
||||
assert resources[0].title == "Dataset1"
|
||||
assert resources[0].description == "desc1"
|
||||
assert resources[1].uri == "rag://dataset/456"
|
||||
assert resources[1].title == "Dataset2"
|
||||
assert resources[1].description == "desc2"
|
||||
|
||||
|
||||
@patch("src.rag.ragflow.requests.get")
|
||||
def test_list_resources_error(mock_get, monkeypatch):
|
||||
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
|
||||
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
|
||||
provider = RAGFlowProvider()
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 500
|
||||
mock_response.text = "fail"
|
||||
mock_get.return_value = mock_response
|
||||
with pytest.raises(Exception):
|
||||
provider.list_resources()
|
||||
@@ -0,0 +1,73 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import pytest
|
||||
|
||||
from src.rag.retriever import Chunk, Document, Resource, Retriever
|
||||
|
||||
|
||||
def test_chunk_init():
|
||||
chunk = Chunk(content="test content", similarity=0.9)
|
||||
assert chunk.content == "test content"
|
||||
assert chunk.similarity == 0.9
|
||||
|
||||
|
||||
def test_document_init_and_to_dict():
|
||||
chunk1 = Chunk(content="chunk1", similarity=0.8)
|
||||
chunk2 = Chunk(content="chunk2", similarity=0.7)
|
||||
doc = Document(
|
||||
id="doc1", url="http://example.com", title="Title", chunks=[chunk1, chunk2]
|
||||
)
|
||||
assert doc.id == "doc1"
|
||||
assert doc.url == "http://example.com"
|
||||
assert doc.title == "Title"
|
||||
assert doc.chunks == [chunk1, chunk2]
|
||||
d = doc.to_dict()
|
||||
assert d["id"] == "doc1"
|
||||
assert d["content"] == "chunk1\n\nchunk2"
|
||||
assert d["url"] == "http://example.com"
|
||||
assert d["title"] == "Title"
|
||||
|
||||
|
||||
def test_document_to_dict_optional_fields():
|
||||
chunk = Chunk(content="only chunk", similarity=1.0)
|
||||
doc = Document(id="doc2", chunks=[chunk])
|
||||
d = doc.to_dict()
|
||||
assert d["id"] == "doc2"
|
||||
assert d["content"] == "only chunk"
|
||||
assert "url" not in d
|
||||
assert "title" not in d
|
||||
|
||||
|
||||
def test_resource_model():
|
||||
resource = Resource(uri="uri1", title="Resource Title")
|
||||
assert resource.uri == "uri1"
|
||||
assert resource.title == "Resource Title"
|
||||
assert resource.description == ""
|
||||
|
||||
|
||||
def test_resource_model_with_description():
|
||||
resource = Resource(uri="uri2", title="Resource2", description="desc")
|
||||
assert resource.description == "desc"
|
||||
|
||||
|
||||
def test_retriever_abstract_methods():
|
||||
class DummyRetriever(Retriever):
|
||||
def list_resources(self, query=None):
|
||||
return [Resource(uri="uri", title="title")]
|
||||
|
||||
def query_relevant_documents(self, query, resources=[]):
|
||||
return [Document(id="id", chunks=[])]
|
||||
|
||||
retriever = DummyRetriever()
|
||||
resources = retriever.list_resources()
|
||||
assert isinstance(resources, list)
|
||||
assert isinstance(resources[0], Resource)
|
||||
docs = retriever.query_relevant_documents("query", resources)
|
||||
assert isinstance(docs, list)
|
||||
assert isinstance(docs[0], Document)
|
||||
|
||||
|
||||
def test_retriever_cannot_instantiate():
|
||||
with pytest.raises(TypeError):
|
||||
Retriever()
|
||||
@@ -0,0 +1,540 @@
|
||||
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.rag.vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider, parse_uri
|
||||
|
||||
|
||||
# Dummy classes to mock dependencies
|
||||
class MockResource:
|
||||
def __init__(self, uri, title="", description=""):
|
||||
self.uri = uri
|
||||
self.title = title
|
||||
self.description = description
|
||||
|
||||
|
||||
class MockChunk:
|
||||
def __init__(self, content, similarity):
|
||||
self.content = content
|
||||
self.similarity = similarity
|
||||
|
||||
|
||||
class MockDocument:
|
||||
def __init__(self, id, title, chunks=None):
|
||||
self.id = id
|
||||
self.title = title
|
||||
self.chunks = chunks or []
|
||||
|
||||
|
||||
# Patch the imports to use mock classes
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_imports():
|
||||
with (
|
||||
patch("src.rag.vikingdb_knowledge_base.Resource", MockResource),
|
||||
patch("src.rag.vikingdb_knowledge_base.Chunk", MockChunk),
|
||||
patch("src.rag.vikingdb_knowledge_base.Document", MockDocument),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def env_vars():
|
||||
"""Fixture to set up environment variables"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE": "10",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_REGION": "cn-north-1",
|
||||
},
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestParseUri:
|
||||
def test_parse_uri_valid_with_fragment(self):
|
||||
"""Test parsing valid URI with fragment"""
|
||||
uri = "rag://dataset/123#doc456"
|
||||
resource_id, document_id = parse_uri(uri)
|
||||
assert resource_id == "123"
|
||||
assert document_id == "doc456"
|
||||
|
||||
def test_parse_uri_valid_without_fragment(self):
|
||||
"""Test parsing valid URI without fragment"""
|
||||
uri = "rag://dataset/123"
|
||||
resource_id, document_id = parse_uri(uri)
|
||||
assert resource_id == "123"
|
||||
assert document_id == ""
|
||||
|
||||
def test_parse_uri_invalid_scheme(self):
|
||||
"""Test parsing URI with invalid scheme"""
|
||||
with pytest.raises(ValueError, match="Invalid URI"):
|
||||
parse_uri("http://dataset/123#abc")
|
||||
|
||||
def test_parse_uri_malformed(self):
|
||||
"""Test parsing malformed URI"""
|
||||
with pytest.raises(ValueError, match="Invalid URI"):
|
||||
parse_uri("invalid_uri")
|
||||
|
||||
|
||||
class TestVikingDBKnowledgeBaseProviderInit:
|
||||
def test_init_success_with_all_env_vars(self, env_vars):
|
||||
"""Test successful initialization with all environment variables"""
|
||||
provider = VikingDBKnowledgeBaseProvider()
|
||||
assert provider.api_url == "api-test.example.com"
|
||||
assert provider.api_ak == "test_ak"
|
||||
assert provider.api_sk == "test_sk"
|
||||
assert provider.retrieval_size == 10
|
||||
assert provider.region == "cn-north-1"
|
||||
assert provider.service == "air"
|
||||
|
||||
def test_init_success_without_retrieval_size(self):
|
||||
"""Test initialization without VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE (should use default)"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
|
||||
},
|
||||
clear=True,
|
||||
):
|
||||
provider = VikingDBKnowledgeBaseProvider()
|
||||
assert provider.retrieval_size == 10
|
||||
|
||||
def test_init_custom_retrieval_size(self):
|
||||
"""Test initialization with custom retrieval size"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE": "5",
|
||||
},
|
||||
):
|
||||
provider = VikingDBKnowledgeBaseProvider()
|
||||
assert provider.retrieval_size == 5
|
||||
|
||||
def test_init_custom_region(self):
|
||||
"""Test initialization with custom region"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_REGION": "us-east-1",
|
||||
},
|
||||
):
|
||||
provider = VikingDBKnowledgeBaseProvider()
|
||||
assert provider.region == "us-east-1"
|
||||
|
||||
def test_init_missing_api_url(self):
|
||||
"""Test initialization fails when API URL is missing"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
|
||||
},
|
||||
clear=True,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_URL is not set"
|
||||
):
|
||||
VikingDBKnowledgeBaseProvider()
|
||||
|
||||
def test_init_missing_api_ak(self):
|
||||
"""Test initialization fails when API AK is missing"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
|
||||
},
|
||||
clear=True,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_AK is not set"
|
||||
):
|
||||
VikingDBKnowledgeBaseProvider()
|
||||
|
||||
def test_init_missing_api_sk(self):
|
||||
"""Test initialization fails when API SK is missing"""
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
|
||||
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
|
||||
},
|
||||
clear=True,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_SK is not set"
|
||||
):
|
||||
VikingDBKnowledgeBaseProvider()
|
||||
|
||||
|
||||
class TestVikingDBKnowledgeBaseProviderSignature:
|
||||
@pytest.fixture
|
||||
def provider(self, env_vars):
|
||||
return VikingDBKnowledgeBaseProvider()
|
||||
|
||||
def test_hmac_sha256(self, provider):
|
||||
"""Test HMAC SHA256 calculation"""
|
||||
key = b"test_key"
|
||||
content = "test_content"
|
||||
result = provider._hmac_sha256(key, content)
|
||||
expected = hmac.new(key, content.encode("utf-8"), hashlib.sha256).digest()
|
||||
assert result == expected
|
||||
|
||||
def test_hash_sha256(self, provider):
|
||||
"""Test SHA256 hash calculation"""
|
||||
data = b"test_data"
|
||||
result = provider._hash_sha256(data)
|
||||
expected = hashlib.sha256(data).digest()
|
||||
assert result == expected
|
||||
|
||||
def test_get_signed_key(self, provider):
|
||||
"""Test signed key generation"""
|
||||
secret_key = "test_secret"
|
||||
date = "20250722"
|
||||
region = "cn-north-1"
|
||||
service = "air"
|
||||
|
||||
result = provider._get_signed_key(secret_key, date, region, service)
|
||||
assert isinstance(result, bytes)
|
||||
assert len(result) == 32 # SHA256 digest is 32 bytes
|
||||
|
||||
def test_create_canonical_request(self, provider):
|
||||
"""Test canonical request creation"""
|
||||
method = "POST"
|
||||
path = "/api/test"
|
||||
query_params = {"param1": "value1", "param2": "value2"}
|
||||
headers = {"Content-Type": "application/json", "Host": "example.com"}
|
||||
payload = b'{"test": "data"}'
|
||||
|
||||
canonical_request, signed_headers = provider._create_canonical_request(
|
||||
method, path, query_params, headers, payload
|
||||
)
|
||||
|
||||
assert "POST" in canonical_request
|
||||
assert "/api/test" in canonical_request
|
||||
assert "param1=value1¶m2=value2" in canonical_request
|
||||
assert "content-type:application/json" in canonical_request
|
||||
assert "host:example.com" in canonical_request
|
||||
assert signed_headers == "content-type;host"
|
||||
|
||||
@patch("src.rag.vikingdb_knowledge_base.datetime")
|
||||
def test_create_signature(self, mock_datetime, provider):
|
||||
"""Test signature creation"""
|
||||
# Mock datetime
|
||||
mock_now = datetime(2025, 7, 22, 10, 30, 45)
|
||||
mock_datetime.utcnow.return_value = mock_now
|
||||
|
||||
method = "POST"
|
||||
path = "/api/test"
|
||||
query_params = {}
|
||||
headers = {}
|
||||
payload = b'{"test": "data"}'
|
||||
|
||||
result = provider._create_signature(
|
||||
method, path, query_params, headers, payload
|
||||
)
|
||||
|
||||
assert "X-Date" in result
|
||||
assert "Host" in result
|
||||
assert "X-Content-Sha256" in result
|
||||
assert "Content-Type" in result
|
||||
assert "Authorization" in result
|
||||
assert "HMAC-SHA256" in result["Authorization"]
|
||||
|
||||
@patch("src.rag.vikingdb_knowledge_base.requests.request")
|
||||
def test_make_signed_request_success(self, mock_request, provider):
|
||||
"""Test successful signed request"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"code": 0, "data": {}}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
result = provider._make_signed_request(
|
||||
"POST", "/api/test", data={"test": "data"}
|
||||
)
|
||||
|
||||
assert result == mock_response
|
||||
mock_request.assert_called_once()
|
||||
|
||||
# Verify the call arguments
|
||||
call_args = mock_request.call_args
|
||||
assert call_args[1]["method"] == "POST"
|
||||
assert call_args[1]["url"] == f"https://{provider.api_url}/api/test"
|
||||
assert call_args[1]["timeout"] == 30
|
||||
|
||||
@patch("src.rag.vikingdb_knowledge_base.requests.request")
|
||||
def test_make_signed_request_with_exception(self, mock_request, provider):
|
||||
"""Test signed request with exception"""
|
||||
mock_request.side_effect = Exception("Network error")
|
||||
|
||||
with pytest.raises(ValueError, match="Request failed: Network error"):
|
||||
provider._make_signed_request("GET", "/api/test")
|
||||
|
||||
|
||||
class TestVikingDBKnowledgeBaseProviderQueryRelevantDocuments:
|
||||
@pytest.fixture
|
||||
def provider(self, env_vars):
|
||||
return VikingDBKnowledgeBaseProvider()
|
||||
|
||||
def test_query_relevant_documents_empty_resources(self, provider):
|
||||
"""Test querying with empty resources list"""
|
||||
result = provider.query_relevant_documents("test query", [])
|
||||
assert result == []
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_query_relevant_documents_success(self, mock_request, provider):
|
||||
"""Test successful document query"""
|
||||
# Mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"result_list": [
|
||||
{
|
||||
"doc_info": {
|
||||
"doc_id": "doc123",
|
||||
"doc_name": "Test Document",
|
||||
},
|
||||
"content": "Test content",
|
||||
"score": 0.95,
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
resources = [MockResource("rag://dataset/123")]
|
||||
result = provider.query_relevant_documents("test query", resources)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].id == "doc123"
|
||||
assert result[0].title == "Test Document"
|
||||
assert len(result[0].chunks) == 1
|
||||
assert result[0].chunks[0].content == "Test content"
|
||||
assert result[0].chunks[0].similarity == 0.95
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_query_relevant_documents_with_document_filter(
|
||||
self, mock_request, provider
|
||||
):
|
||||
"""Test document query with document ID filter"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"code": 0, "data": {"result_list": []}}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
resources = [MockResource("rag://dataset/123#doc456")]
|
||||
provider.query_relevant_documents("test query", resources)
|
||||
|
||||
# Verify that query_param with doc_filter was included in the request
|
||||
call_args = mock_request.call_args
|
||||
request_data = call_args[1]["data"]
|
||||
assert "query_param" in request_data
|
||||
assert "doc_filter" in request_data["query_param"]
|
||||
|
||||
doc_filter = request_data["query_param"]["doc_filter"]
|
||||
assert doc_filter["op"] == "must"
|
||||
assert doc_filter["field"] == "doc_id"
|
||||
assert doc_filter["conds"] == ["doc456"]
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_query_relevant_documents_api_error(self, mock_request, provider):
|
||||
"""Test handling of API error response"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"code": 1, "message": "API Error"}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
resources = [MockResource("rag://dataset/123")]
|
||||
with pytest.raises(
|
||||
ValueError, match="Failed to query documents from resource: API Error"
|
||||
):
|
||||
provider.query_relevant_documents("test query", resources)
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_query_relevant_documents_json_decode_error(self, mock_request, provider):
|
||||
"""Test handling of JSON decode error"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
resources = [MockResource("rag://dataset/123")]
|
||||
with pytest.raises(ValueError, match="Failed to parse JSON response"):
|
||||
provider.query_relevant_documents("test query", resources)
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_query_relevant_documents_multiple_resources(self, mock_request, provider):
|
||||
"""Test querying multiple resources and merging results"""
|
||||
# Mock responses for different resources
|
||||
responses = [
|
||||
{
|
||||
"code": 0,
|
||||
"data": {
|
||||
"result_list": [
|
||||
{
|
||||
"doc_info": {
|
||||
"doc_id": "doc1",
|
||||
"doc_name": "Document 1",
|
||||
},
|
||||
"content": "Content 1",
|
||||
"score": 0.9,
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
{
|
||||
"code": 0,
|
||||
"data": {
|
||||
"result_list": [
|
||||
{
|
||||
"doc_info": {
|
||||
"doc_id": "doc1",
|
||||
"doc_name": "Document 1",
|
||||
},
|
||||
"content": "Content 2",
|
||||
"score": 0.8,
|
||||
},
|
||||
{
|
||||
"doc_info": {
|
||||
"doc_id": "doc2",
|
||||
"doc_name": "Document 2",
|
||||
},
|
||||
"content": "Content 3",
|
||||
"score": 0.7,
|
||||
},
|
||||
]
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
mock_responses = [MagicMock() for _ in responses]
|
||||
for i, resp in enumerate(responses):
|
||||
mock_responses[i].json.return_value = resp
|
||||
mock_request.side_effect = mock_responses
|
||||
|
||||
resources = [
|
||||
MockResource("rag://dataset/123"),
|
||||
MockResource("rag://dataset/456"),
|
||||
]
|
||||
result = provider.query_relevant_documents("test query", resources)
|
||||
|
||||
# Should have 2 documents: doc1 (with 2 chunks) and doc2 (with 1 chunk)
|
||||
assert len(result) == 2
|
||||
doc1 = next(doc for doc in result if doc.id == "doc1")
|
||||
doc2 = next(doc for doc in result if doc.id == "doc2")
|
||||
assert len(doc1.chunks) == 2
|
||||
assert len(doc2.chunks) == 1
|
||||
|
||||
|
||||
class TestVikingDBKnowledgeBaseProviderListResources:
|
||||
@pytest.fixture
|
||||
def provider(self, env_vars):
|
||||
return VikingDBKnowledgeBaseProvider()
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_list_resources_success(self, mock_request, provider):
|
||||
"""Test successful resource listing"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"collection_list": [
|
||||
{
|
||||
"resource_id": "123",
|
||||
"collection_name": "Dataset 1",
|
||||
"description": "Description 1",
|
||||
},
|
||||
{
|
||||
"resource_id": "456",
|
||||
"collection_name": "Dataset 2",
|
||||
"description": "Description 2",
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
result = provider.list_resources()
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0].uri == "rag://dataset/123"
|
||||
assert result[0].title == "Dataset 1"
|
||||
assert result[0].description == "Description 1"
|
||||
assert result[1].uri == "rag://dataset/456"
|
||||
assert result[1].title == "Dataset 2"
|
||||
assert result[1].description == "Description 2"
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_list_resources_with_query_filter(self, mock_request, provider):
|
||||
"""Test resource listing with query filter"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
"code": 0,
|
||||
"data": {
|
||||
"collection_list": [
|
||||
{
|
||||
"resource_id": "123",
|
||||
"collection_name": "Test Dataset",
|
||||
"description": "Description",
|
||||
},
|
||||
{
|
||||
"resource_id": "456",
|
||||
"collection_name": "Other Dataset",
|
||||
"description": "Description",
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
result = provider.list_resources("test")
|
||||
|
||||
# Should only return the dataset with "test" in the name
|
||||
assert len(result) == 1
|
||||
assert result[0].title == "Test Dataset"
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_list_resources_api_error(self, mock_request, provider):
|
||||
"""Test handling of API error in list_resources"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"code": 1, "message": "API Error"}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
with pytest.raises(Exception, match="Failed to list resources: API Error"):
|
||||
provider.list_resources()
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_list_resources_json_decode_error(self, mock_request, provider):
|
||||
"""Test handling of JSON decode error in list_resources"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
with pytest.raises(ValueError, match="Failed to parse JSON response"):
|
||||
provider.list_resources()
|
||||
|
||||
@patch.object(VikingDBKnowledgeBaseProvider, "_make_signed_request")
|
||||
def test_list_resources_empty_response(self, mock_request, provider):
|
||||
"""Test handling of empty response"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {"code": 0, "data": {"collection_list": []}}
|
||||
mock_request.return_value = mock_response
|
||||
|
||||
result = provider.list_resources()
|
||||
assert result == []
|
||||
Reference in New Issue
Block a user