wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add router session mining tools

Browse Source
This commit is contained in:
武阳
2026-04-30 18:25:13 +08:00
parent 97bf5b9a54
commit 2ea3a63447
5 changed files with 1240 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_tools.py
+229
View File
@@ -35,6 +35,13 @@ from .data_agent_inputs import (
load_input_sources,
render_source_context,
)
from .data_agent_router_sessions import (
DataAgentRouterSessionError,
convert_router_candidates_to_records,
profile_router_sessions,
sample_router_candidates,
search_router_sessions,
)
from .session_env_vars import get_session_env_vars
if TYPE_CHECKING:
@@ -1466,6 +1473,137 @@ def default_tool_registry() -> dict[str, AgentTool]:
},
handler=_data_agent_render_source_context_tool,
),
AgentTool(
name='data_agent_profile_router_sessions',
description=(
'Profile local router_session_parquet date partitions before mining online sessions. '
'Returns schema, sampled row count, distributions, and example sessions.'
),
parameters={
'type': 'object',
'properties': {
'dates': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Date partitions such as ["20260428"]. Resolved under router_session_parquet/date=YYYYMMDD.',
},
'paths': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Workspace-relative parquet files or date directories. Use when dates is not enough.',
},
'max_files': {'type': 'integer', 'minimum': 1, 'maximum': 200},
'max_rows_per_file': {'type': 'integer', 'minimum': 1, 'maximum': 100000},
'examples_limit': {'type': 'integer', 'minimum': 0, 'maximum': 20},
},
},
handler=_data_agent_profile_router_sessions_tool,
),
AgentTool(
name='data_agent_search_router_sessions',
description=(
'Search local router_session_parquet data with reviewable filters such as device, domain, intent, func, '
'query keywords, regex, turn count, and match scope. Returns turn-level candidates with prev turns.'
),
parameters={
'type': 'object',
'properties': {
'dates': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Date partitions such as ["20260428"]. Resolved under router_session_parquet/date=YYYYMMDD.',
},
'paths': {
'type': 'array',
'items': {'type': 'string'},
'description': 'Workspace-relative parquet files or date directories. Use when dates is not enough.',
},
'devices': {'type': 'array', 'items': {'type': 'string'}},
'domains': {'type': 'array', 'items': {'type': 'string'}},
'intents': {'type': 'array', 'items': {'type': 'string'}},
'funcs': {'type': 'array', 'items': {'type': 'string'}},
'query_keywords': {'type': 'array', 'items': {'type': 'string'}},
'keyword_match_mode': {'type': 'string', 'enum': ['any', 'all']},
'query_regex': {'type': 'string'},
'match_scope': {'type': 'string', 'enum': ['any_turn', 'last_turn']},
'turn_count_min': {'type': 'integer', 'minimum': 1},
'turn_count_max': {'type': 'integer', 'minimum': 1},
'max_files': {'type': 'integer', 'minimum': 1, 'maximum': 200},
'max_rows_per_file': {'type': 'integer', 'minimum': 1, 'maximum': 100000},
'max_candidates': {'type': 'integer', 'minimum': 1, 'maximum': 5000},
},
},
handler=_data_agent_search_router_sessions_tool,
),
AgentTool(
name='data_agent_sample_router_candidates',
description=(
'Sample router session candidates for human review and return candidate-level summary statistics. '
'Use after data_agent_search_router_sessions.'
),
parameters={
'type': 'object',
'properties': {
'candidates': {
'oneOf': [
{'type': 'array'},
{'type': 'object'},
{'type': 'string'},
],
'description': 'Candidates array, search result object, or JSON string containing candidates.',
},
'sample_size': {'type': 'integer', 'minimum': 1, 'maximum': 1000},
'strategy': {'type': 'string', 'enum': ['random', 'first', 'stride']},
'seed': {'type': 'integer'},
},
'required': ['candidates'],
},
handler=_data_agent_sample_router_candidates_tool,
),
AgentTool(
name='data_agent_convert_router_candidates_to_records',
description=(
'Convert reviewed online router session candidates directly into canonical data-agent records. '
'Use this when the user wants to keep mined online data as samples; do not use generation-plan tools for this path.'
),
parameters={
'type': 'object',
'properties': {
'candidates': {
'oneOf': [
{'type': 'array'},
{'type': 'object'},
{'type': 'string'},
],
'description': 'Candidates array, search/sample result object, or JSON string containing candidates.',
},
'dataset_label': {'type': 'string'},
'default_target': {
'type': 'string',
'description': 'Target label used for included candidates unless review_decisions provides a target.',
},
'review_decisions': {
'type': 'array',
'description': 'Optional include/exclude/uncertain decisions keyed by semantic_session_id or req_id.',
'items': {
'type': 'object',
'properties': {
'semantic_session_id': {'type': 'string'},
'req_id': {'type': 'string'},
'matched_turn_index': {'type': 'integer'},
'decision': {'type': 'string', 'enum': ['include', 'exclude', 'uncertain']},
'target': {'type': 'string'},
'notes': {'type': 'string'},
},
},
},
'batch_id': {'type': 'string'},
'include_uncertain': {'type': 'boolean'},
},
'required': ['candidates', 'dataset_label'],
},
handler=_data_agent_convert_router_candidates_to_records_tool,
),
AgentTool(
name='data_agent_show_generation_plan',
description='Show a data-agent generation plan for human review.',
@@ -1689,6 +1827,15 @@ def _coerce_int(arguments: dict[str, Any], key: str, default: int) -> int:
return value
def _optional_int(arguments: dict[str, Any], key: str) -> int | None:
value = arguments.get(key)
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
raise ToolExecutionError(f'{key} must be an integer')
return value
def _coerce_float(arguments: dict[str, Any], key: str, default: float) -> float:
value = arguments.get(key, default)
if isinstance(value, bool) or not isinstance(value, (int, float)):
@@ -1855,6 +2002,88 @@ def _data_agent_render_source_context_tool(arguments: dict[str, Any], context: T
return data_agent_dumps_payload(payload)
def _data_agent_profile_router_sessions_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
try:
payload = profile_router_sessions(
context.root,
dates=_optional_string_list(arguments, 'dates'),
paths=_optional_string_list(arguments, 'paths'),
max_files=_coerce_int(arguments, 'max_files', 8),
max_rows_per_file=_coerce_int(arguments, 'max_rows_per_file', 2000),
examples_limit=_coerce_int(arguments, 'examples_limit', 3),
)
except DataAgentRouterSessionError as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_search_router_sessions_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
try:
payload = search_router_sessions(
context.root,
dates=_optional_string_list(arguments, 'dates'),
paths=_optional_string_list(arguments, 'paths'),
devices=_optional_string_list(arguments, 'devices'),
domains=_optional_string_list(arguments, 'domains'),
intents=_optional_string_list(arguments, 'intents'),
funcs=_optional_string_list(arguments, 'funcs'),
query_keywords=_optional_string_list(arguments, 'query_keywords'),
keyword_match_mode=_optional_string(arguments, 'keyword_match_mode') or 'any',
query_regex=_optional_string(arguments, 'query_regex'),
match_scope=_optional_string(arguments, 'match_scope') or 'any_turn',
turn_count_min=_optional_int(arguments, 'turn_count_min'),
turn_count_max=_optional_int(arguments, 'turn_count_max'),
max_files=_coerce_int(arguments, 'max_files', 20),
max_rows_per_file=_coerce_int(arguments, 'max_rows_per_file', 5000),
max_candidates=_coerce_int(arguments, 'max_candidates', 200),
)
except DataAgentRouterSessionError as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_sample_router_candidates_tool(arguments: dict[str, Any], _context: ToolExecutionContext) -> str:
if 'candidates' not in arguments:
raise ToolExecutionError('candidates is required')
try:
payload = sample_router_candidates(
arguments['candidates'],
sample_size=_coerce_int(arguments, 'sample_size', 100),
strategy=_optional_string(arguments, 'strategy') or 'random',
seed=_coerce_int(arguments, 'seed', 20260430),
)
except DataAgentRouterSessionError as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_convert_router_candidates_to_records_tool(
arguments: dict[str, Any],
_context: ToolExecutionContext,
) -> str:
if 'candidates' not in arguments:
raise ToolExecutionError('candidates is required')
review_decisions = arguments.get('review_decisions')
if review_decisions is not None and not isinstance(review_decisions, list):
raise ToolExecutionError('review_decisions must be an array')
include_uncertain = arguments.get('include_uncertain', False)
if not isinstance(include_uncertain, bool):
raise ToolExecutionError('include_uncertain must be a boolean')
try:
payload = convert_router_candidates_to_records(
arguments['candidates'],
dataset_label=_require_string(arguments, 'dataset_label'),
default_target=_optional_string(arguments, 'default_target'),
review_decisions=review_decisions,
batch_id=_optional_string(arguments, 'batch_id') or 'router',
include_uncertain=include_uncertain,
)
payload['validation'] = validate_dataset_records(payload['records'])
except (DataAgentRouterSessionError, DataRecordError) as exc:
raise ToolExecutionError(str(exc)) from exc
return data_agent_dumps_payload(payload)
def _data_agent_show_generation_plan_tool(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
try:
payload = get_generation_plan(