Add router session mining tools
This commit is contained in:
武阳
@@ -0,0 +1,229 @@
|
||||
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_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()
|
||||
Reference in New Issue
Block a user