武阳
253 lines
9.9 KiB
Python
253 lines
9.9 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
|
|
from src.agent_types import AgentRuntimeConfig
|
|
from src.data_agent_router_sessions import (
|
|
convert_router_candidates_to_records,
|
|
profile_router_sessions,
|
|
sample_router_candidates,
|
|
search_router_sessions,
|
|
)
|
|
|
|
|
|
HAS_PYARROW = importlib.util.find_spec('pyarrow') is not None
|
|
|
|
|
|
class DataAgentRouterSessionTests(unittest.TestCase):
|
|
def test_sample_router_candidates_accepts_search_payload(self) -> None:
|
|
payload = {
|
|
'candidates': [
|
|
_candidate('s1', 'phone', 'QA', 'unknown', '怎么查找设备'),
|
|
_candidate('s2', 'car', 'mapCopilot', '导航', '导航去公司'),
|
|
_candidate('s3', 'phone', 'QA', 'unknown', '手机丢了怎么办'),
|
|
]
|
|
}
|
|
|
|
result = sample_router_candidates(payload, sample_size=2, strategy='stride')
|
|
|
|
self.assertEqual(result['total_candidate_count'], 3)
|
|
self.assertEqual(result['sampled_count'], 2)
|
|
self.assertEqual(result['summary']['device'][0], {'value': 'phone', 'count': 2})
|
|
self.assertEqual(result['candidates'][0]['semantic_session_id'], 's1')
|
|
self.assertEqual(result['candidates'][1]['semantic_session_id'], 's3')
|
|
|
|
def test_sample_router_candidates_accepts_json_string(self) -> None:
|
|
result = sample_router_candidates(
|
|
json.dumps([_candidate('s1', 'phone', 'QA', 'unknown', '怎么查找设备')], ensure_ascii=False),
|
|
sample_size=1,
|
|
)
|
|
|
|
self.assertEqual(result['sampled_count'], 1)
|
|
self.assertEqual(result['candidates'][0]['matched_turn']['query'], '怎么查找设备')
|
|
|
|
def test_convert_router_candidates_to_records_keeps_online_metadata(self) -> None:
|
|
candidate = _candidate('s1', 'phone', 'QA', 'unknown', '总结一下')
|
|
candidate['turn_count'] = 2
|
|
candidate['matched_turn_index'] = 1
|
|
candidate['matched_turn']['timestamp'] = 20
|
|
candidate['prev_turns'] = [{'timestamp': 10, 'query': '这篇文章讲了什么', 'tts': ''}]
|
|
|
|
result = convert_router_candidates_to_records(
|
|
[candidate],
|
|
dataset_label='总结类边界评测集',
|
|
default_target='Summarize',
|
|
batch_id='summary',
|
|
)
|
|
|
|
self.assertEqual(result['record_count'], 1)
|
|
record = result['records'][0]
|
|
self.assertEqual(record['source']['type'], 'online')
|
|
self.assertEqual(record['source']['request_id'], 's1-rid')
|
|
self.assertEqual(record['turn']['query'], '总结一下')
|
|
self.assertEqual(record['prev_session'], [{'query': '这篇文章讲了什么', 'tts': '', 'timestamp': 10}])
|
|
self.assertEqual(record['context']['domain'], 'QA')
|
|
self.assertEqual(record['label']['target'], 'Summarize')
|
|
|
|
def test_convert_router_candidates_to_records_applies_review_decisions(self) -> None:
|
|
candidates = [
|
|
_candidate('s1', 'phone', 'QA', 'unknown', '总结一下'),
|
|
_candidate('s2', 'phone', 'QA', 'unknown', '总结一下红楼梦'),
|
|
]
|
|
|
|
result = convert_router_candidates_to_records(
|
|
candidates,
|
|
dataset_label='总结类边界评测集',
|
|
review_decisions=[
|
|
{
|
|
'semantic_session_id': 's1',
|
|
'matched_turn_index': 0,
|
|
'decision': 'include',
|
|
'target': 'Summarize',
|
|
'notes': '前文有可总结内容',
|
|
},
|
|
{
|
|
'semantic_session_id': 's2',
|
|
'matched_turn_index': 0,
|
|
'decision': 'exclude',
|
|
'notes': '开放知识对象',
|
|
},
|
|
],
|
|
)
|
|
|
|
self.assertEqual(result['record_count'], 1)
|
|
self.assertEqual(result['skipped_count'], 1)
|
|
self.assertEqual(result['records'][0]['meta']['notes'], '前文有可总结内容')
|
|
|
|
@unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests')
|
|
def test_profile_and_search_router_sessions_read_parquet(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
_write_router_parquet(Path(tmp_dir))
|
|
|
|
profile = profile_router_sessions(tmp_dir, dates=['20260428'], max_files=1)
|
|
search = search_router_sessions(
|
|
tmp_dir,
|
|
dates=['20260428'],
|
|
domains=['mapCopilot'],
|
|
query_keywords=['导航'],
|
|
max_files=1,
|
|
max_candidates=5,
|
|
)
|
|
|
|
self.assertEqual(profile['sampled_row_count'], 2)
|
|
self.assertEqual(profile['missing_req_id_count'], 0)
|
|
self.assertTrue(any(item['value'] == 'phone' for item in profile['distributions']['device']))
|
|
self.assertEqual(search['candidate_count'], 1)
|
|
candidate = search['candidates'][0]
|
|
self.assertEqual(candidate['req_id'], 'rid-1')
|
|
self.assertEqual(candidate['matched_turn']['query'], '导航去公司')
|
|
self.assertEqual(candidate['prev_turns'][0]['query'], '你好小爱')
|
|
|
|
@unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests')
|
|
def test_router_session_tools_allow_external_absolute_paths(self) -> None:
|
|
with tempfile.TemporaryDirectory() as workspace_dir, tempfile.TemporaryDirectory() as data_dir:
|
|
_write_router_parquet(Path(data_dir))
|
|
external_partition = Path(data_dir) / 'router_session_parquet' / 'date=20260428'
|
|
|
|
profile = profile_router_sessions(
|
|
workspace_dir,
|
|
paths=[str(external_partition)],
|
|
max_files=1,
|
|
)
|
|
search = search_router_sessions(
|
|
workspace_dir,
|
|
paths=[str(external_partition)],
|
|
query_keywords=['导航'],
|
|
max_files=1,
|
|
max_candidates=5,
|
|
)
|
|
|
|
self.assertEqual(profile['sampled_row_count'], 2)
|
|
self.assertEqual(search['candidate_count'], 1)
|
|
self.assertEqual(search['candidates'][0]['matched_turn']['query'], '导航去公司')
|
|
|
|
@unittest.skipUnless(HAS_PYARROW, 'pyarrow is required for parquet tests')
|
|
def test_router_session_tools_execute_against_registry(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
_write_router_parquet(Path(tmp_dir))
|
|
context = build_tool_context(AgentRuntimeConfig(cwd=Path(tmp_dir)))
|
|
registry = default_tool_registry()
|
|
|
|
search_result = execute_tool(
|
|
registry,
|
|
'data_agent_search_router_sessions',
|
|
{'dates': ['20260428'], 'query_regex': '查找设备|手机丢了', 'max_files': 1},
|
|
context,
|
|
)
|
|
self.assertTrue(search_result.ok, search_result.content)
|
|
sample_result = execute_tool(
|
|
registry,
|
|
'data_agent_sample_router_candidates',
|
|
{'candidates': json.loads(search_result.content), 'sample_size': 1, 'strategy': 'first'},
|
|
context,
|
|
)
|
|
convert_result = execute_tool(
|
|
registry,
|
|
'data_agent_convert_router_candidates_to_records',
|
|
{
|
|
'candidates': json.loads(sample_result.content),
|
|
'dataset_label': '查找设备评测集',
|
|
'default_target': 'QA',
|
|
'batch_id': 'find_device',
|
|
},
|
|
context,
|
|
)
|
|
|
|
self.assertTrue(sample_result.ok, sample_result.content)
|
|
self.assertEqual(json.loads(sample_result.content)['sampled_count'], 1)
|
|
self.assertTrue(convert_result.ok, convert_result.content)
|
|
self.assertTrue(json.loads(convert_result.content)['validation']['ok'])
|
|
|
|
|
|
def _candidate(session_id: str, device: str, domain: str, intent: str, query: str) -> dict[str, object]:
|
|
return {
|
|
'semantic_session_id': session_id,
|
|
'req_id': f'{session_id}-rid',
|
|
'device': device,
|
|
'turn_count': 1,
|
|
'matched_turn_index': 0,
|
|
'matched_turn': {
|
|
'timestamp': 1,
|
|
'query': query,
|
|
'tts': '',
|
|
'action': {'domain': domain, 'intent': intent},
|
|
},
|
|
'prev_turns': [],
|
|
'source_path': 'router_session_parquet/date=20260428/part-0.parquet',
|
|
'date': '20260428',
|
|
}
|
|
|
|
|
|
def _write_router_parquet(root: Path) -> None:
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
|
|
partition = root / 'router_session_parquet' / 'date=20260428'
|
|
partition.mkdir(parents=True)
|
|
rows = [
|
|
{
|
|
'semantic_session_id': 'phone_20260428_1',
|
|
'device': 'phone',
|
|
'req_id': 'rid-1',
|
|
'turn_count': 2,
|
|
'turns_array': [
|
|
{
|
|
'timestamp': 1,
|
|
'query': '你好小爱',
|
|
'action_json': '{"domain":"dialogCopilot","intent":"问候"}',
|
|
'tts': '你好',
|
|
},
|
|
{
|
|
'timestamp': 2,
|
|
'query': '导航去公司',
|
|
'action_json': '{"domain":"mapCopilot","intent":"导航","func":"start_navigation"}',
|
|
'tts': '开始导航',
|
|
},
|
|
],
|
|
},
|
|
{
|
|
'semantic_session_id': 'phone_20260428_2',
|
|
'device': 'phone',
|
|
'req_id': 'rid-2',
|
|
'turn_count': 1,
|
|
'turns_array': [
|
|
{
|
|
'timestamp': 3,
|
|
'query': '手机丢了怎么查找设备',
|
|
'action_json': '{"domain":"QA","intent":"unknown"}',
|
|
'tts': '',
|
|
}
|
|
],
|
|
},
|
|
]
|
|
pq.write_table(pa.Table.from_pylist(rows), partition / 'part-0.parquet')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|