mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-08 20:34:48 +08:00
c960dc2a4c
### What problem does this PR solve? Refine handling of POST /api/v1/datasets/search in GO ### Type of change - [x] Refactoring
595 lines
28 KiB
Python
595 lines
28 KiB
Python
#
|
||
#
|
||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
#
|
||
"""
|
||
Search-datasets consistency tests between Python (port 9380) and Go (port 9384) servers.
|
||
Compares /api/v1/datasets/search endpoint responses for consistency between Python and Go.
|
||
|
||
When an LLM is involved (rerank_id, keyword, or cross_languages is set), both sides
|
||
call the LLM independently and the chunk *count*, ordering, and scores can drift
|
||
across runs and between Python/Go. The test logs the per-chunk chunk_id + similarity
|
||
for human inspection but skips strict count/order/score comparisons in those cases.
|
||
|
||
When no LLM is involved, both servers run the same deterministic retrieval path
|
||
with no non-deterministic dependencies, so the responses are expected to be
|
||
byte-identical: same chunk count, same chunk order, and the same per-field
|
||
values (with the empty-value normalization done in compare_chunks).
|
||
|
||
All datasets and documents are created once at module level, then each test unit
|
||
runs against the pre-built data. Cleanup happens automatically at module teardown.
|
||
"""
|
||
import logging
|
||
import os
|
||
import sys
|
||
import pytest
|
||
import requests
|
||
import tempfile
|
||
import time
|
||
import uuid
|
||
|
||
# Logging setup
|
||
# Default is silent. Set LOG_LEVEL=INFO locally to see logger.info() messages.
|
||
_LOG_LEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
|
||
_log_handler = logging.StreamHandler(sys.stderr)
|
||
_log_handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
|
||
logger = logging.getLogger("search_datasets_consistency")
|
||
logger.setLevel(getattr(logging, _LOG_LEVEL, logging.WARNING))
|
||
logger.addHandler(_log_handler)
|
||
logger.propagate = False
|
||
|
||
from test.testcases.utils import wait_for
|
||
|
||
PYTHON_HOST = "http://localhost:9380"
|
||
GO_HOST = "http://localhost:9384"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Shared test documents
|
||
# ---------------------------------------------------------------------------
|
||
THREE_KINGDOMS_TXT = """
|
||
曹操（155 年—220 年）
|
||
作为曹魏政权的奠基者，曹操展现出卓越的政治、军事和文学才能。在政治上，
|
||
他挟天子以令诸侯，稳固自身政治地位，推行一系列政策，如推行屯田制，不仅
|
||
解决了粮食短缺问题，还使大量流民得以安置，促进了农业生产的恢复与发展 ；
|
||
在军事方面，曹操一生征战无数，官渡之战以少胜多击败袁绍，统一北方，压制
|
||
匈奴等异族势力，为北方地区带来相对稳定的局面 。然而，曹操的评价褒贬不
|
||
一，"治世能臣"体现他在治理国家、施展政治抱负方面的才能；"乱世奸雄"
|
||
则反映他在乱世中为达目的不择手段的一面 。
|
||
|
||
刘备（161 年—223 年）
|
||
蜀汉开国皇帝刘备，以"仁德"著称。他的一生充满传奇色彩，早年与关羽、张
|
||
飞桃园结义，奠定兄弟情谊，三人患难与共，为兴复汉室而努力 。刘备求贤若
|
||
渴，三顾茅庐请出诸葛亮，得到这位智谋之士的辅佐，为蜀汉政权的建立与发展
|
||
奠定根基 。刘备善于笼络人心，凭借自身的人格魅力吸引了众多人才，在乱世
|
||
中逐渐崛起，建立蜀汉，与曹魏、东吴形成三足鼎立之势 。
|
||
|
||
孙权（182 年—252 年）
|
||
东吴的建立者孙权，年少继承父兄基业，展现出非凡的领导才能 。他擅长用人
|
||
制衡，麾下人才济济，周瑜、鲁肃、陆逊等皆是东吴的栋梁之才 。孙权依托长
|
||
江天险，制定了稳健的战略，致力于开发江南经济与海外贸易，使东吴在三国中
|
||
占据重要地位 。在赤壁之战中，孙权与刘备联军击败曹操，奠定三国鼎立的基
|
||
础；之后又在夷陵之战中击败刘备，巩固了东吴在江南的统治 。
|
||
|
||
赵云（？—229 年）
|
||
赵云一生未败，长坂坡单骑救主，在曹操大军中七进七出，如入无人之境，成功
|
||
救出刘备之子刘禅，被誉为"常胜将军" 。赵云跟随刘备多年，忠心耿耿，多
|
||
次在危难时刻挺身而出，保护刘备及其家人的安全 。他武艺高强，为人正直，
|
||
深受刘备和蜀军将士的敬重 。
|
||
"""
|
||
|
||
WATER_MARGIN_TXT = """
|
||
宋江（？—？）
|
||
宋江是梁山泊一百零八将之首，人称"及时雨"。他原为郓城县押司，因仗义疏财、
|
||
济人贫苦，在江湖上广结英雄豪杰。宋江性格矛盾，一方面讲究忠义，一心想着
|
||
招安报国；另一方面又领导梁山好汉对抗朝廷。他带领梁山好汉两赢童贯、三败
|
||
高俅，最终接受朝廷招安，率军征讨辽国、田虎、王庆、方腊，立下赫赫战功。
|
||
|
||
武松（？—？）
|
||
武松因其排行第二，人称"武二郎"。景阳冈打虎，使他一举成名。武松武艺高强，
|
||
力大无穷，是梁山泊步军头领。他因兄长武大郎被西门庆、潘金莲毒害，怒杀二
|
||
人，被发配孟州。在孟州，他醉打蒋门神，帮助施恩夺回快活林。大闹飞云浦、
|
||
血溅鸳鸯楼后，武松走上反抗道路，最终加入梁山。
|
||
|
||
李逵（？—？）
|
||
李逵是梁山泊好汉中最具反抗精神的代表，人称"黑旋风"。他性格鲁莽直率，
|
||
嫉恶如仇，对宋江忠心耿耿。李逵使用两把板斧，在战场上所向披靡，屡立战功。
|
||
他曾在江州劫法场救宋江，在沂岭杀四虎，展现了惊人的勇气和力量。
|
||
"""
|
||
|
||
ENGLISH_DOCS = """
|
||
Artificial Intelligence
|
||
Artificial intelligence (AI) is intelligence demonstrated by machines, in contrast to
|
||
the natural intelligence displayed by humans and animals. Leading AI textbooks define
|
||
the field as the study of "intelligent agents": any system that perceives its environment
|
||
and takes actions that maximize its chance of successfully achieving its goals.
|
||
AI applications include advanced web search engines (e.g., Google Search), recommendation
|
||
systems (used by YouTube, Amazon, and Netflix), understanding human speech (such as Siri
|
||
and Alexa), self-driving cars (e.g., Waymo), generative and creative tools (ChatGPT and
|
||
AI art), and superhuman play and analysis in strategy games (such as chess and Go).
|
||
|
||
Machine Learning
|
||
Machine learning (ML) is a field of inquiry devoted to understanding and building methods
|
||
that "learn" – that is, methods that leverage data to improve performance on some set of
|
||
tasks. It is seen as a part of artificial intelligence. Machine learning algorithms build
|
||
a model based on sample data, known as training data, in order to make predictions or
|
||
decisions without being explicitly programmed to do so. Machine learning algorithms are
|
||
used in a wide variety of applications, such as in medicine, email filtering, speech
|
||
recognition, agriculture, and computer vision, where it is difficult or unfeasible to
|
||
develop conventional algorithms to perform the needed tasks.
|
||
|
||
Deep Learning
|
||
Deep learning is part of a broader family of machine learning methods based on artificial
|
||
neural networks with representation learning. Learning can be supervised, semi-supervised
|
||
or unsupervised. Deep-learning architectures such as deep neural networks, deep belief
|
||
networks, deep reinforcement learning, recurrent neural networks, convolutional neural
|
||
networks and transformers have been applied to fields including computer vision, speech
|
||
recognition, natural language processing, machine translation, bioinformatics, drug design,
|
||
medical image analysis, climate science, material inspection and board game programs,
|
||
where they have produced results comparable to and in some cases surpassing human expert
|
||
performance.
|
||
"""
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
def compare_chunks(python_chunk, go_chunk):
|
||
"""Compare a single chunk between Python and Go responses."""
|
||
all_keys = set(python_chunk.keys()) | set(go_chunk.keys())
|
||
|
||
for field in all_keys:
|
||
p_val = python_chunk.get(field)
|
||
g_val = go_chunk.get(field)
|
||
|
||
# Normalize empty values to None so the comparison treats
|
||
# "no value" shapes interchangeably across engines:
|
||
# Python Infinity: [] (keyword list, infinity_conn.py:793-794)
|
||
# Python Elasticsearch: "" (keyword, es_conn.py:626-644)
|
||
# Go (both engines): None (intentional — see retrieval.go:403-424)
|
||
if isinstance(p_val, (list, str)) and not p_val:
|
||
p_val = None
|
||
if isinstance(g_val, (list, str)) and not g_val:
|
||
g_val = None
|
||
|
||
if p_val is None or g_val is None:
|
||
if p_val != g_val:
|
||
raise AssertionError(f"Field '{field}': python={p_val}, go={g_val}")
|
||
continue
|
||
|
||
if field in ("similarity", "term_similarity", "vector_similarity"):
|
||
if p_val != g_val:
|
||
raise AssertionError(
|
||
f"Field '{field}' differs: python={p_val}, go={g_val}, diff={abs(p_val - g_val)}"
|
||
)
|
||
elif isinstance(p_val, (list, dict)):
|
||
if p_val != g_val:
|
||
raise AssertionError(f"Field '{field}' mismatch")
|
||
else:
|
||
if p_val != g_val:
|
||
raise AssertionError(f"Field '{field}': python={p_val}, go={g_val}")
|
||
|
||
|
||
def search_and_compare(rest_client, dataset_ids, cfg):
|
||
"""Perform search on both servers and compare results.
|
||
|
||
dataset_ids can be a single dataset ID string or a list of dataset IDs.
|
||
cfg is a dict with keys matching the search payload:
|
||
question (required), top_k (default 5), rerank_id, search_id, keyword,
|
||
vector_similarity_weight, similarity_threshold, use_kg, cross_languages,
|
||
page, size, meta_data_filter
|
||
"""
|
||
headers = {"Authorization": f"Bearer {rest_client.token}"}
|
||
if isinstance(dataset_ids, str):
|
||
ids = [dataset_ids]
|
||
else:
|
||
ids = list(dataset_ids)
|
||
|
||
search_payload = {
|
||
"dataset_ids": ids,
|
||
"question": cfg["question"],
|
||
"top_k": cfg.get("top_k", 5),
|
||
}
|
||
optional_fields = [
|
||
"rerank_id", "search_id", "keyword", "vector_similarity_weight",
|
||
"similarity_threshold", "use_kg", "cross_languages", "page", "size",
|
||
"meta_data_filter", "doc_ids",
|
||
]
|
||
for field in optional_fields:
|
||
value = cfg.get(field)
|
||
if value is not None:
|
||
search_payload[field] = value
|
||
|
||
# Call Python server
|
||
python_res = rest_client.post("/datasets/search", json=search_payload)
|
||
assert python_res.status_code == 200, f"Python server error: {python_res.status_code}, body: {python_res.text}"
|
||
python_data = python_res.json()
|
||
assert python_data["code"] == 0, f"Python payload error: {python_data}"
|
||
|
||
# Call Go server with same auth
|
||
go_res = requests.post(
|
||
f"{GO_HOST}/api/v1/datasets/search",
|
||
json=search_payload,
|
||
headers=headers,
|
||
timeout=30,
|
||
)
|
||
assert go_res.status_code == 200, f"Go server error: {go_res.status_code}, body: {go_res.text}"
|
||
go_data = go_res.json()
|
||
assert go_data["code"] == 0, f"Go payload error: {go_data}"
|
||
|
||
python_chunks = python_data["data"]["chunks"]
|
||
go_chunks = go_data["data"]["chunks"]
|
||
|
||
logger.info(f"python_chunks={len(python_chunks)}, go_chunks={len(go_chunks)}")
|
||
logger.info(f" Python chunks: {[(c.get('chunk_id','?'), c.get('similarity',0)) for c in python_chunks]}")
|
||
logger.info(f" Go chunks: {[(c.get('chunk_id','?'), c.get('similarity',0)) for c in go_chunks]}")
|
||
|
||
llm_involved = bool(cfg.get("rerank_id") or cfg.get("keyword") or cfg.get("cross_languages"))
|
||
if not llm_involved:
|
||
assert len(python_chunks) == len(go_chunks), \
|
||
f"Chunk count differs: python={len(python_chunks)}, go={len(go_chunks)}"
|
||
for i, (p_chunk, g_chunk) in enumerate(zip(python_chunks, go_chunks)):
|
||
try:
|
||
compare_chunks(p_chunk, g_chunk)
|
||
except AssertionError as e:
|
||
raise AssertionError(f"Chunk {i} comparison failed: {e}")
|
||
python_total = python_data["data"].get("total", 0)
|
||
go_total = go_data["data"].get("total", 0)
|
||
if python_total != go_total:
|
||
raise AssertionError(f"total differs: python={python_total}, go={go_total}")
|
||
|
||
return len(python_chunks), len(python_chunks)
|
||
|
||
|
||
def _upload_and_parse(rest_client, dataset_id, text, filename="doc.txt"):
|
||
"""Upload text as a file and wait for parsing to complete. Returns document_id."""
|
||
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as f:
|
||
f.write(text)
|
||
temp_path = f.name
|
||
|
||
with open(temp_path, 'rb') as f:
|
||
files = [("file", (filename, f))]
|
||
upload_res = rest_client.post(f"/datasets/{dataset_id}/documents", files=files)
|
||
assert upload_res.status_code == 200, f"Failed to upload {filename}: {upload_res.text}"
|
||
doc_id = upload_res.json()["data"][0]["id"]
|
||
|
||
parse_res = rest_client.post(
|
||
f"/datasets/{dataset_id}/documents/parse",
|
||
json={"document_ids": [doc_id]},
|
||
)
|
||
assert parse_res.status_code == 200, f"Failed to start parsing {filename}: {parse_res.text}"
|
||
|
||
@wait_for(120, 2, f"Document parsing timeout for {filename}")
|
||
def check_parsed():
|
||
doc_res = rest_client.get(f"/datasets/{dataset_id}/documents", params={"id": doc_id})
|
||
if doc_res.status_code != 200:
|
||
return False
|
||
docs = doc_res.json()["data"]["docs"]
|
||
return bool(docs) and docs[0].get("run") == "DONE"
|
||
|
||
check_parsed()
|
||
return doc_id
|
||
|
||
|
||
def _get_chunk_count(rest_client, dataset_id, doc_id):
|
||
"""Return the number of chunks for a parsed document."""
|
||
doc_res = rest_client.get(f"/datasets/{dataset_id}/documents", params={"id": doc_id})
|
||
docs = doc_res.json()["data"]["docs"]
|
||
return docs[0].get("chunk_count", 0)
|
||
|
||
|
||
def _set_metadata(rest_client, dataset_id, doc_id, updates):
|
||
"""Set metadata key-value pairs on a document."""
|
||
headers = {"Authorization": f"Bearer {rest_client.token}"}
|
||
formatted_updates = [{"key": k, "value": v} for k, v in updates.items()]
|
||
res = requests.patch(
|
||
f"{PYTHON_HOST}/api/v1/datasets/{dataset_id}/documents/metadatas",
|
||
json={
|
||
"selector": {"document_ids": [doc_id]},
|
||
"updates": formatted_updates,
|
||
},
|
||
headers=headers,
|
||
timeout=30,
|
||
)
|
||
assert res.status_code == 200, f"Failed to set metadata: {res.text}"
|
||
assert res.json()["code"] == 0, f"Metadata update error: {res.json()}"
|
||
meta_str = ", ".join(f"{k}={v}" for k, v in updates.items())
|
||
logger.info(f" metadata set on doc {doc_id}: {meta_str}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Module-level fixture: create all datasets once, delete at the end
|
||
# ---------------------------------------------------------------------------
|
||
@pytest.fixture(scope="module")
|
||
def all_datasets(rest_client):
|
||
"""Set up all datasets and documents once for the module.
|
||
|
||
Returns a dict with:
|
||
ds_chinese : 1 dataset with 2 files (Three Kingdoms + Water Margin, with metadata)
|
||
ds_chinese_doc1 : doc_id for Three Kingdoms
|
||
ds_chinese_doc2 : doc_id for Water Margin
|
||
|
||
ds_chinese_2 : 1 dataset with Three Kingdoms only
|
||
ds_3k_doc : doc_id
|
||
|
||
ds_english : 1 dataset with English text
|
||
ds_english_doc : doc_id
|
||
"""
|
||
logger.info("\n[SETUP] Creating all datasets for search-datasets consistency tests...")
|
||
|
||
data = {}
|
||
|
||
# -----------------------------------------------------------------------
|
||
# 1) 1 dataset with 2 files (Chinese)
|
||
# -----------------------------------------------------------------------
|
||
create_res = rest_client.post("/datasets", json={
|
||
"name": "consistency_chinese",
|
||
"embedding_model": "BAAI/bge-small-en-v1.5@Builtin",
|
||
"parser_config": {"chunk_token_num": 1, "delimiter": "`\n\n`"},
|
||
})
|
||
assert create_res.status_code == 200, create_res.text
|
||
assert create_res.json()["code"] == 0, create_res.json()
|
||
ds_chinese_id = create_res.json()["data"]["id"]
|
||
|
||
doc1 = _upload_and_parse(rest_client, ds_chinese_id, THREE_KINGDOMS_TXT, "three_kingdoms.txt")
|
||
doc2 = _upload_and_parse(rest_client, ds_chinese_id, WATER_MARGIN_TXT, "water_margin.txt")
|
||
logger.info(f" ds_chinese: {ds_chinese_id} (3k={doc1}, wm={doc2})")
|
||
logger.info(f" 3K chunks: {_get_chunk_count(rest_client, ds_chinese_id, doc1)}")
|
||
logger.info(f" WM chunks: {_get_chunk_count(rest_client, ds_chinese_id, doc2)}")
|
||
|
||
# Set metadata on individual documents
|
||
_set_metadata(rest_client, ds_chinese_id, doc1, {"era": 220, "source": "luo", "character": ["曹操", "刘备", "孙权", "赵云"]})
|
||
_set_metadata(rest_client, ds_chinese_id, doc2, {"era": 960, "source": "shi", "character": ["宋江", "武松", "李逵"]})
|
||
|
||
data["ds_chinese"] = ds_chinese_id
|
||
data["ds_chinese_doc1"] = doc1
|
||
data["ds_chinese_doc2"] = doc2
|
||
|
||
# -----------------------------------------------------------------------
|
||
# 2) 1 dataset with Three Kingdoms only
|
||
# -----------------------------------------------------------------------
|
||
create_res = rest_client.post("/datasets", json={
|
||
"name": "consistency_three_kingdoms",
|
||
"embedding_model": "BAAI/bge-small-en-v1.5@Builtin",
|
||
"parser_config": {"chunk_token_num": 1, "delimiter": "`\n\n`"},
|
||
})
|
||
assert create_res.status_code == 200, create_res.text
|
||
assert create_res.json()["code"] == 0, create_res.json()
|
||
ds_3k_id = create_res.json()["data"]["id"]
|
||
|
||
doc_3k = _upload_and_parse(rest_client, ds_3k_id, THREE_KINGDOMS_TXT, "three_kingdoms.txt")
|
||
logger.info(f" ds_chinese_2: {ds_3k_id} (doc={doc_3k})")
|
||
logger.info(f" chunks: {_get_chunk_count(rest_client, ds_3k_id, doc_3k)}")
|
||
|
||
data["ds_chinese_2"] = ds_3k_id
|
||
data["ds_3k_doc"] = doc_3k
|
||
|
||
# Wait for metadata to be indexed
|
||
time.sleep(2)
|
||
|
||
# -----------------------------------------------------------------------
|
||
# 3) 1 dataset with English text
|
||
# -----------------------------------------------------------------------
|
||
create_res = rest_client.post("/datasets", json={
|
||
"name": "consistency_english",
|
||
"embedding_model": "BAAI/bge-small-en-v1.5@Builtin",
|
||
"parser_config": {"chunk_token_num": 1, "delimiter": "`\n\n`"},
|
||
})
|
||
assert create_res.status_code == 200, create_res.text
|
||
assert create_res.json()["code"] == 0, create_res.json()
|
||
ds_en_id = create_res.json()["data"]["id"]
|
||
|
||
doc_en = _upload_and_parse(rest_client, ds_en_id, ENGLISH_DOCS, "english_docs.txt")
|
||
logger.info(f" ds_english: {ds_en_id} (doc={doc_en})")
|
||
logger.info(f" chunks: {_get_chunk_count(rest_client, ds_en_id, doc_en)}")
|
||
|
||
data["ds_english"] = ds_en_id
|
||
data["ds_english_doc"] = doc_en
|
||
|
||
logger.info("[SETUP] All datasets ready.\n")
|
||
|
||
yield data
|
||
|
||
# Teardown: delete all datasets
|
||
logger.info("\n[TEARDOWN] Deleting all datasets...")
|
||
all_ids = [
|
||
data["ds_chinese"],
|
||
data["ds_chinese_2"],
|
||
data["ds_english"],
|
||
]
|
||
res = rest_client.delete("/datasets", json={"ids": all_ids})
|
||
assert res.status_code == 200, f"Teardown failed: {res.text}"
|
||
logger.info("[TEARDOWN] Done.")
|
||
|
||
|
||
# Skip every test in this module from CI. Remove the next line to re-enable.
|
||
pytestmark = pytest.mark.skipif(
|
||
os.getenv("CI") == "true",
|
||
reason="GO server is not started in CI",
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test Unit 1: Search consistency — 1 dataset with 2 files
|
||
# ---------------------------------------------------------------------------
|
||
@pytest.mark.p2
|
||
def test_search_datasets_consistency_basic(rest_client, all_datasets):
|
||
"""
|
||
Compare /api/v1/datasets/search responses between Python and Go servers for consistency.
|
||
Tests the single dataset (ds_chinese) which contains 2 files.
|
||
"""
|
||
dataset_id = all_datasets["ds_chinese"]
|
||
doc1 = all_datasets["ds_chinese_doc1"]
|
||
doc2 = all_datasets["ds_chinese_doc2"]
|
||
logger.info(f"Using dataset (Chinese, 2 files): {dataset_id} (doc1={doc1}, doc2={doc2})")
|
||
|
||
search_configs = [
|
||
{"question": "曹操"},
|
||
{"question": "曹操", "page": 2, "size": 2},
|
||
{"question": "曹操", "top_k": 2},
|
||
{"question": "曹操", "similarity_threshold": 0.0},
|
||
{"question": "曹操", "similarity_threshold": 0.5},
|
||
{"question": "曹操", "vector_similarity_weight": 0.0},
|
||
{"question": "曹操", "vector_similarity_weight": 0.7},
|
||
{"question": "曹操", "keyword": True},
|
||
{"question": "努力发展农业", "keyword": True},
|
||
{"question": "political status", "cross_languages": ["Chinese"]},
|
||
{"question": "诸葛亮"},
|
||
{"question": "努力发展农业"},
|
||
{"question": "曹操 诸葛亮 周瑜"},
|
||
{"question": "曹操", "top_k": 3, "rerank_id": "BAAI/bge-reranker-v2-m3@CI@SILICONFLOW"},
|
||
{"question": "曹操", "doc_ids": [doc1]},
|
||
{"question": "曹操", "doc_ids": [doc2]},
|
||
{"question": "曹操", "doc_ids": []},
|
||
]
|
||
|
||
for cfg in search_configs:
|
||
cfg_str = ", ".join(f"{k}={v}" for k, v in cfg.items())
|
||
logger.info(f"\n--- Testing: {cfg_str} ---")
|
||
total, chunk_count = search_and_compare(rest_client, dataset_id, cfg)
|
||
logger.info(f"SUCCESS: Python and Go responses match for {chunk_count} chunks, total={total}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test Unit 2: Metadata filter consistency
|
||
# ---------------------------------------------------------------------------
|
||
@pytest.mark.p2
|
||
def test_search_datasets_consistency_metadata_filter(rest_client, all_datasets):
|
||
"""
|
||
Compare Python vs Go search with metadata filtering on the Chinese dataset for consistency.
|
||
Uses ds_chinese which has 2 documents with metadata:
|
||
- doc1: era=220, source=luo, character=[曹操,刘备,孙权,赵云]
|
||
- doc2: era=960, source=shi, character=[宋江,武松,李逵]
|
||
"""
|
||
dataset_id = all_datasets["ds_chinese"]
|
||
logger.info(f"Using dataset (Chinese, metadata): {dataset_id}")
|
||
|
||
search_configs = [
|
||
# Manual filters
|
||
{"question": "打虎", "meta_data_filter": {"method": "manual", "manual": [{"key": "era", "op": "=", "value": 960}]}},
|
||
{"question": "曹操", "meta_data_filter": {"method": "manual", "manual": [{"key": "source", "op": "=", "value": "luo"}]}},
|
||
{"question": "打虎", "meta_data_filter": {"method": "manual", "manual": [{"key": "era", "op": "≠", "value": 960}]}},
|
||
{"question": "打虎", "meta_data_filter": {"method": "manual", "manual": [{"key": "era", "op": ">", "value": 220}]}},
|
||
{"question": "曹操", "meta_data_filter": {"method": "manual", "manual": [{"key": "source", "op": "contains", "value": "luo"}]}},
|
||
{"question": "努力发展农业", "meta_data_filter": {"method": "manual", "manual": [{"key": "character", "op": "in", "value": ["曹操","孙权"]}]}},
|
||
{"question": "打虎", "meta_data_filter": {"method": "manual", "manual": [{"key": "character", "op": "=", "value": "武松"}]}},
|
||
]
|
||
|
||
for cfg in search_configs:
|
||
cfg_str = ", ".join(f"{k}={v}" for k, v in cfg.items())
|
||
logger.info(f"\n--- Metadata filter: {cfg_str} ---")
|
||
total, chunk_count = search_and_compare(rest_client, dataset_id, cfg)
|
||
logger.info(f"SUCCESS: {chunk_count} chunks, total={total}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test Unit 3: Search with search_id
|
||
# ---------------------------------------------------------------------------
|
||
@pytest.mark.p2
|
||
def test_search_datasets_consistency_with_search_id(rest_client, all_datasets):
|
||
"""
|
||
Compare Python vs Go with search_id parameter for consistency.
|
||
Creates a search config with multiple parameters set, then tests that
|
||
both servers honor the stored config overrides.
|
||
"""
|
||
dataset_id = all_datasets["ds_chinese"]
|
||
logger.info(f"Using dataset (Chinese): {dataset_id}")
|
||
|
||
# Create a search config first
|
||
search_name = f"consistency_search_{uuid.uuid4().hex[:8]}"
|
||
search_res = rest_client.post("/searches", json={"name": search_name, "description": "consistency test search"})
|
||
assert search_res.status_code == 200, f"Failed to create search: {search_res.text}"
|
||
search_payload = search_res.json()
|
||
assert search_payload["code"] == 0, f"Search creation error: {search_payload}"
|
||
search_id = search_payload["data"]["search_id"]
|
||
logger.info(f"Created search_id: {search_id}")
|
||
|
||
# Update the search config with multiple parameters
|
||
search_config = {
|
||
"similarity_threshold": 0.2,
|
||
"vector_similarity_weight": 0.7,
|
||
"top_k": 3,
|
||
"use_kg": False,
|
||
"keyword": False,
|
||
}
|
||
update_res = rest_client.put(
|
||
f"/searches/{search_id}",
|
||
json={"name": search_name, "search_config": search_config},
|
||
)
|
||
assert update_res.status_code == 200, f"Failed to update search config: {update_res.text}"
|
||
assert update_res.json()["code"] == 0, f"Search config update error: {update_res.json()}"
|
||
logger.info(f"Updated search config: {search_config}")
|
||
|
||
# Search with search_id — config overrides should apply
|
||
search_configs = [
|
||
{"question": "曹操", "search_id": search_id},
|
||
{"question": "曹操 诸葛亮", "search_id": search_id},
|
||
]
|
||
|
||
for cfg in search_configs:
|
||
cfg_str = ", ".join(f"{k}={v}" for k, v in cfg.items())
|
||
logger.info(f"\n--- Testing with search_id: {cfg_str} ---")
|
||
total, chunk_count = search_and_compare(rest_client, dataset_id, cfg)
|
||
logger.info(f"SUCCESS: Python and Go responses match for {chunk_count} chunks, total={total}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test Unit 4: Multi-dataset search
|
||
# ---------------------------------------------------------------------------
|
||
@pytest.mark.p2
|
||
def test_search_datasets_consistency_multi_dataset(rest_client, all_datasets):
|
||
"""
|
||
Compare Python vs Go when searching across 2 datasets simultaneously for consistency.
|
||
Uses ds_chinese and ds_chinese_2.
|
||
"""
|
||
both_ids = [all_datasets["ds_chinese"], all_datasets["ds_chinese_2"]]
|
||
logger.info(f"Using datasets (multi): ds_chinese={all_datasets['ds_chinese']}, ds_chinese_2={all_datasets['ds_chinese_2']}")
|
||
|
||
search_configs = [
|
||
{"question": "武松打虎"},
|
||
{"question": "努力发展农业"},
|
||
{"question": "曹操 宋江"},
|
||
]
|
||
|
||
for cfg in search_configs:
|
||
cfg_str = ", ".join(f"{k}={v}" for k, v in cfg.items())
|
||
logger.info(f"\n--- Multi-dataset Testing: {cfg_str} ---")
|
||
total, chunk_count = search_and_compare(rest_client, both_ids, cfg)
|
||
logger.info(f"SUCCESS: Python and Go responses match for {chunk_count} chunks, total={total}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test Unit 5: English search consistency
|
||
# ---------------------------------------------------------------------------
|
||
@pytest.mark.p2
|
||
def test_search_datasets_consistency_english(rest_client, all_datasets):
|
||
"""
|
||
Compare Python vs Go with English text documents for consistency.
|
||
Uses ds_english.
|
||
"""
|
||
dataset_id = all_datasets["ds_english"]
|
||
logger.info(f"Using English dataset: {dataset_id}")
|
||
|
||
search_configs = [
|
||
{"question": "artificial intelligence"},
|
||
{"question": "neural networks and deep learning"},
|
||
{"question": "neural networks", "keyword": True},
|
||
{"question": "convolutional neural networks", "keyword": True},
|
||
{"question": "人工智能", "cross_languages": ["English"]},
|
||
{"question": "机器学习", "cross_languages": ["English"]},
|
||
]
|
||
|
||
for cfg in search_configs:
|
||
cfg_str = ", ".join(f"{k}={v}" for k, v in cfg.items())
|
||
logger.info(f"\n--- Testing (English): {cfg_str} ---")
|
||
total, chunk_count = search_and_compare(rest_client, dataset_id, cfg)
|
||
logger.info(f"SUCCESS: Python and Go responses match for {chunk_count} chunks, total={total}")
|