--- name: alibabacloud-data-agent-skill description: | Invoke Alibaba Cloud Apsara Data Agent for Analytics via CLI to perform natural language-driven data analysis on enterprise databases. Data Agent for Analytics is an intelligent data analysis agent developed by Alibaba Cloud Database team for enterprise users. It automatically completes requirement analysis, data understanding, analysis insights, and report generation based on natural language descriptions. This tool supports: discovering data resources (instances/databases/tables) managed in DMS, initiating query or deep analysis sessions, real-time progress tracking, and retrieving analysis conclusions and generated reports. Use this Skill when users need to query databases, analyze data trends, generate data reports, ask questions in natural language, or mention "Data Agent", "data analysis", "database query", "SQL analysis", "data insights". compatibility: | Requires valid Alibaba Cloud credentials (default credential chain or API_KEY); Requires dependencies in requirements.txt to be installed; Data sources must be managed in Alibaba Cloud Apsara Database or DMS. domain: AIOps --- metadata: author: DataAgent Team version: "1.7.2" --- # Changelog - **v1.7.2**: Use Alibaba Cloud default credential chain instead of explicit AK/SK, add User-Agent header, fix RAM policy wildcard issues - **v1.7.1**: Fix CLI `ls` command API response parsing (support case-insensitive field names), optimize SKILL documentation structure, separate ANALYSIS mode specification document - **v1.7.0**: API_KEY authentication support, native async execution mode, session isolation, enhanced attach mode, optimized log output --- --- # Installation ## Configure Credentials This Skill uses Alibaba Cloud default credential chain (recommended) or API_KEY authentication. ### Option 1: Default Credential Chain (Recommended) The Skill uses Alibaba Cloud SDK's default credential chain to automatically obtain credentials, supporting environment variables, configuration files, instance roles, etc. See [Alibaba Cloud Credential Chain Documentation](https://help.aliyun.com/document_detail/378659.html) ### Option 2: API_KEY Authentication (File Analysis Only) ```bash export DATA_AGENT_API_KEY=your-api-key export DATA_AGENT_REGION=cn-hangzhou ``` Get API_KEY: [Data Agent Console](https://agent.dms.aliyun.com/cn-hangzhou/api-key) ### Permission Requirements RAM users need `AliyunDMSFullAccess` or `AliyunDMSDataAgentFullAccess` permissions. See [RAM-POLICIES.md](references/RAM-POLICIES.md) for detailed permission information. ## Debug Mode ```bash DATA_AGENT_DEBUG_API=1 python3 scripts/data_agent_cli.py file example.csv -q "analyze" ``` ## 💡 Getting Started Tips - Use the built-in demo database `internal_data_employees` (DataAgent's built-in test database containing employee, department, and salary data) for first-time experience - Or use local file `assets/example_game_data.csv` for file analysis experience # Data Agent CLI — Unified Command-Line Data Analysis Tool ## Overview `scripts/data_agent_cli.py` helps users complete the full workflow from **discover data → initiate analysis → track progress → get results**. ### Core Concepts > **⚠️ Key Prerequisite**: Data Agent can only analyze databases that have been **imported into Data Agent Data Center**. > > - **Data Center**: Data Agent's data center, only databases here can be analyzed > - **DMS**: Alibaba Cloud Data Management Service, stores metadata of all databases > - **Relationship**: Databases registered in DMS ≠ Databases in Data Center > > **Usage Flow**: > 1. First use `ls` to check if the target database exists in Data Center > 2. If **not found**, use `dms` subcommand to search for database info, then use `import` subcommand to import it > 3. After successful import, you can use `db` subcommand for analysis --- ## Analysis Modes - **ASK_DATA** (default): Synchronous execution, sub-second response, suitable for quick Q&A - **ANALYSIS**: Deep analysis, takes 5-40 minutes, requires spawning a sub-agent for async execution or using --async-run parameter > See [ANALYSIS_MODE.md](references/ANALYSIS_MODE.md) for details --- ## Session Reuse Use `db`/`file` to create a session for initial analysis, then use `attach --session-id ` to reuse the session for follow-up questions. > See [COMMANDS.md](references/COMMANDS.md) and [WORKFLOWS.md](references/WORKFLOWS.md) for details --- ## Quick Start ```bash # 1. List available databases python3 scripts/data_agent_cli.py ls # 2. Query analysis (synchronous response) python3 scripts/data_agent_cli.py db \ --dms-instance-id --dms-db-id \ --instance-name --db-name \ --tables "employees,departments" -q "Which department has the highest average salary" # 3. Follow-up question (reuse session) python3 scripts/data_agent_cli.py attach --session-id -q "Break down by month" ``` > 📖 See [WORKFLOWS.md](references/WORKFLOWS.md) and [COMMANDS.md](references/COMMANDS.md) for complete workflows, command reference, and best practices --- ## Project Structure ``` # Skill root directory ├── SKILL.md # This document ├── scripts/ # Source code │ ├── data_agent/ # SDK module │ ├── cli/ # CLI module │ ├── data_agent_cli.py # CLI entry point │ └── requirements.txt # Dependencies ├── sessions/ # Session data └── references/ # Reference documents ``` FILE:assets/HEARTBEAT.md ## Data Agent 自动监控当 HEARTBEAT 触发时，按照以下步骤 chain-of-thought 执行： 1. 扫描目录：`dms-data-agent/sessions/*`（或你的实际 sessions 路径） - 使用 file_list 或 ls 列出所有子目录（每个对应一个 session_id） 2. 对每个 session 检查状态： - 读 `status.txt` - 如果 "running"： - 读 `progress.log` 把新的进展告诉用户，尤其是计划执行的步骤内容 - 如果 "completed"： - 读 reports/ 目录或使用report子命令获取总结文件， Markdown 格式的内容给客户呈现， html可以通过浏览器打开。 - 如果 "failed"： - 读 error.log，总结原因 - 建议："请用 attach {id} 重试或检查日志" 3. 通知规则： - 需要及时的反馈阶段性洞察给客户，以及里程碑首次到达、完成或失败时发消息。 - 使用 send_message tool 或当前 channel（如 Telegram/WhatsApp）推送 - 避免重复：记录上次计划执行的步骤结论（可写到 SESSION-STATE.md 或专用 data-agent-state.md） 4. 如果无新通知内容： - 安静结束 turn（回复 HEARTBEAT_OK，gateway 会自动丢弃）优先使用 isolated agentTurn 执行检查（不干扰主对话）。 FILE:references/ANALYSIS_MODE.md # ANALYSIS Mode Sub-Agent Specification ANALYSIS mode tasks are long-running (5-40 minutes), **you must spawn a sub-agent to execute in background** to avoid blocking the main conversation flow. --- ## Main Agent Behavior 1. Inform user that deep analysis has started, provide Session ID 2. Return immediately, continue responding to other user questions (non-blocking) --- ## Sub-Agent Behavior 1. Execute command in background, periodically poll `status.txt` (every 30-60 seconds) 2. Push progress updates to main agent/user at key checkpoints 3. Return report or guide user to console when task completes --- ## Key Checkpoints for Sub-Agent to Report | Checkpoint | Trigger Condition | Report Content | Required Action | |------|----------|----------|----------| | 🚀 Task Started | After command execution | Session ID, estimated duration | No action needed | | 📋 Execution Plan Ready | `status.txt` = `waiting_input` | Show analysis plan/SQL | **Must wait for user confirmation, then `attach -q 'confirm execution'`** | | 📊 Phase Conclusion | New content in `progress.log` | Phase analysis summary | No action needed | | ✅ Task Completed | `status.txt` = `completed` | Report link or file path | Guide user to view report | | ❌ Task Failed | `status.txt` = `failed` | Error message | Prompt user to retry | > ⚠️ **`waiting_input` is a mandatory wait point**: Worker has exited, sub-agent must present execution plan to user, execute `attach -q 'confirm execution'` after receiving confirmation, otherwise task will be permanently paused. --- ## Guide User to View Report After Task Completion ```bash # Download report to local python3 scripts/data_agent_cli.py reports --session-id # Or guide user to Data Agent console (recommended) # https://agent.dms.aliyun.com//session/ ``` FILE:references/COMMANDS.md # Subcommand Reference This document contains detailed descriptions and complete parameter lists for all subcommands. --- ## ls Subcommand — Discover Data Resources Before initiating analysis, use `ls` to understand what databases and tables are available. > **⚠️ Important**: The `ls` command only shows databases that have been **imported into Data Agent Data Center**. If you can't find the database you need, it means the database has not been imported from DMS yet. ### List All Databases ```bash python3 scripts/data_agent_cli.py ls ``` Output is divided into two groups: - **Database Connections** (ImportType: RDS/DMS) — Real relational databases, can be used with `db` subcommand - **File Data Sources** (ImportType: FILE) — Uploaded file datasets Each database record displays: ``` internal_data_employees [mysql] (RDS) AgentDbId : DmsDbId : DmsInstanceId : InstanceName : ``` > **Tip**: `internal_data_employees` is DataAgent's built-in demo database, containing employee, department, and salary test data, suitable for first-time experience. ### Filter by Keyword ```bash python3 scripts/data_agent_cli.py ls --search internal_data_employees ``` ### List Tables for a Specific Database + Generate Usable Command ```bash python3 scripts/data_agent_cli.py ls --db-id ``` Output includes a `db` command template ready to copy and use. --- ## db Subcommand — Database Analysis > **⚠️ Important**: The `db` subcommand requires the database to **already exist in Data Agent Data Center**. ### Data Source Parameters | Parameter | Description | |------|------| | `--dms-instance-id` | DMS Instance ID (numeric) | | `--dms-db-id` | DMS Database ID (numeric) | | `--instance-name` | RDS Instance Name | | `--db-name` | Database Name | | `--tables` | Table name list, comma-separated | | `--engine` | Database engine, default `mysql` | ### Session Parameters | Parameter | Description | |------|------| | `--session-mode` | `ASK_DATA` (default) / `ANALYSIS` / `INSIGHT` | | `--output` | `summary` (default) / `detail` / `raw` | | `--enable-search` | Enable search capability (default `false`) | ### Query Methods | Parameter | Description | |------|------| | `-q` / `--query` | Single query (runs preset query if not specified) | --- ## file Subcommand — File Analysis File analysis uses `ANALYSIS` mode by default. ### Parameters | Parameter | Description | |------|------| | `--session-mode` | `ASK_DATA` / `ANALYSIS` (default) / `INSIGHT` | | `--output` | `summary` (default) / `detail` / `raw` | | `--enable-search` | Enable search capability (default `false`) | | `-q` / `--query` | Custom query question | --- ## reports Subcommand — Get Generated Results View and download reports and chart files generated by the session. ### Parameters | Parameter | Description | |------|------| | `--session-id` | **Required**, Session ID | --- ## Session Mode Description | Mode | Features | Duration | Report | Use Cases | |------|------|------|------|----------| | `ASK_DATA` | Quick SQL query + natural language response, supports follow-up questions | 30 sec - 2 min | ❌ | Simple queries, data validation | | `ANALYSIS` | Deep analysis, generates complete report | 5-15 min | ✅ | Topic analysis, business reports | | `INSIGHT` | Multi-dimensional, in-depth data insights | 20-40 min | ✅ | Strategic analysis, trend research | > **⚠️ Important**: When you need to generate reports (HTML/Markdown/charts), **ANALYSIS mode is recommended**. ## dms Subcommand — DMS Tool Integration Direct access to DMS metadata, used for discovering instances, databases, and tables. ### Tool List | Tool | Description | |------|------| | `list-instances` | Search database instance list in DMS | | `search-database` | Search databases by schema name | | `list-tables` | List tables in specified database | --- ## import Subcommand — Import DMS Database Import DMS database tables into Data Agent Data Center. ### Parameters | Parameter | Description | |------|------| | `--dms-instance-id` | **Required**, DMS Instance ID (numeric) | | `--dms-db-id` | **Required**, DMS Database ID (numeric) | | `--instance-name` | **Required**, RDS Instance Name | | `--db-name` | **Required**, Database Name | | `--tables` | **Required**, Table name list to import (comma-separated) | | `--engine` | Database engine type (default: mysql) | | `--region` | Region ID (default: cn-hangzhou) | --- ## attach Subcommand — Session Reuse Connect to an existing session to continue conversation, confirm plans, or check progress. ### Parameters | Parameter | Description | |------|------| | `--session-id` | **Required**, Session ID to connect to | | `-q` / `--query` | Send single query | | `--from-start` | Replay session history from beginning (equivalent to --checkpoint 0) | | `--checkpoint` | Specify exact checkpoint to resume from (e.g., `--checkpoint 219`), used for precise recovery after network interruption | | `--output` | `summary` (default) / `detail` / `raw` | FILE:references/RAM-POLICIES.md # RAM Permission Policies This Skill requires the following Alibaba Cloud RAM permissions: ## Required Permissions | Policy | Description | |----------|------| | `AliyunDMSFullAccess` | Full access to DMS Data Management Service | | `AliyunDMSDataAgentFullAccess` | Full access to Data Agent | ## Minimal Permission Policy (Recommended) If you only need to use Data Agent features, you can configure the following minimal permissions: ```json { "Version": "1", "Statement": [ { "Effect": "Allow", "Action": [ "dms:ListInstances", "dms:ListDatabases", "dms:ListTables", "dms:ListColumns", "dms:GetInstance", "dms:GetDatabase", "dms:GetTableTopology", "dms:GetMetaTableDetailInfo", "dms:DescribeInstance", "dms:CreateDataAgentSession", "dms:DescribeDataAgentSession", "dms:ListDataAgentSession", "dms:SendChatMessage", "dms:GetChatContent", "dms:DescribeDataAgentUsage", "dms:UpdateDataAgentSession", "dms:CreateDataAgentFeedback", "dms:DescribeFileUploadSignature", "dms:FileUploadCallback", "dms:ListFileUpload", "dms:DeleteFileUpload", "dms:ListDataCenterDatabase", "dms:ListDataCenterTable", "dms:AddDataCenterTable" ], "Resource": "*" } ] } ``` ## Configuration Instructions 1. Log in to [Alibaba Cloud RAM Console](https://ram.console.aliyun.com/) 2. Create or select a user 3. Add the above permission policies to the user 4. Create AccessKey for Skill authentication FILE:references/WORKFLOWS.md # Typical Workflows This document contains complete operational workflow examples. --- ## Method 1: Start from Existing Database in Data Center (Recommended) > **Prerequisite**: The database must already exist in Data Agent Data Center. > > **Built-in Demo Database**: `internal_data_employees` is DataAgent's built-in test database, containing employee, department, and salary data, suitable for first-time experience. ``` Step 1 ls -- List available databases Step 2 ls --db-id -- List tables for a specific database, and print db command ready to copy Step 3 db -q -- Initiate query/analysis session, output progress and conclusions in real-time Step 4 attach -- Connect to existing session (confirm plan / follow-up / view latest results) ``` ### Complete Example ```bash # Step 1: Discover databases python3 scripts/data_agent_cli.py ls # Step 2: View tables in built-in test database internal_data_employees and get command template python3 scripts/data_agent_cli.py ls --db-id # Step 3: Query (copy command from previous step, replace question) python3 scripts/data_agent_cli.py db \ --dms-instance-id --dms-db-id \ --instance-name --db-name internal_data_employees \ --tables "employees,departments,salaries" \ --session-mode ASK_DATA \ -q "Which position has the highest salary" # Step 4: Connect to existing session for follow-up questions python3 scripts/data_agent_cli.py attach --session-id -q "Calculate average salary by department" ``` --- ## Method 2: Discover and Import from DMS Instance to Data Center > **Use Case**: When Data Agent Data Center doesn't have the database you need. ``` Step 1 dms list-instances -- Query database instances in DMS Step 2 dms search-database -- Search databases in instance Step 3 dms list-tables -- List tables in database Step 4 ls -- Check if Data Center already has this database Step 5 import -- Import DMS database tables to Data Center Step 6 db -- Initiate query/analysis session Step 7 attach -- Connect to existing session ``` ### Complete Example ```bash # 1. Query DMS instance list python3 scripts/data_agent_cli.py dms list-instances # 2. Search for target database (get Database ID) python3 scripts/data_agent_cli.py dms search-database --search-key employees # 3. View tables in database python3 scripts/data_agent_cli.py dms list-tables --database-id # 4. Check if Data Center already has this database python3 scripts/data_agent_cli.py ls --search employees # 5. Import to Data Center (required step) python3 scripts/data_agent_cli.py import \ --dms-instance-id \ --dms-db-id \ --instance-name \ --db-name employees \ --tables "departments,employees,salaries" # 6. Initiate analysis python3 scripts/data_agent_cli.py db \ --dms-instance-id --dms-db-id \ --instance-name --db-name employees \ --tables "departments,employees,salaries" \ -q "Query the department with highest average salary" ``` --- ## Background Execution Best Practices For long-running `ls` and `db` commands, running in background is recommended: ```bash # Start ANALYSIS task in background nohup python3 scripts/data_agent_cli.py db \ --dms-instance-id --dms-db-id \ --instance-name --db-name internal_data_employees \ --tables "employees,departments,salaries" \ --session-mode ANALYSIS \ -q "Analyze salary distribution and generate report" > analysis.log 2>&1 & # Get session ID from log grep "Session ready" analysis.log # Attach anytime to check progress python3 scripts/data_agent_cli.py attach --session-id # If network is interrupted or you want to resume from a specific state, specify checkpoint python3 scripts/data_agent_cli.py attach --session-id --checkpoint ``` **Benefits**: - Avoid task failure due to network interruption (seamless resume with `--checkpoint` parameter) - Check progress anytime via `attach` - Output logs for later review --- ## Session Reuse Workflow For multiple analyses on the same database, reusing sessions is recommended for efficiency: ```bash # Analysis 1: Create new session python3 scripts/data_agent_cli.py db \ --dms-instance-id --dms-db-id \ --instance-name --db-name internal_data_employees \ --tables "employees,departments" \ --session-mode ANALYSIS \ -q "Analyze 2024 salary growth trends" # Output: ✅ Async task started. Session ID: abc123xyz # Analysis 2: Reuse same session, follow-up with details python3 scripts/data_agent_cli.py attach --session-id abc123xyz -q "Break down salary structure by job level" # Analysis 3: Modify plan python3 scripts/data_agent_cli.py attach --session-id abc123xyz -q "Simplify to 3 steps" # Analysis 4: Confirm execution python3 scripts/data_agent_cli.py attach --session-id abc123xyz -q "Confirm execution" # Step 5: Read final results cat sessions/abc123xyz/progress.log # Step 6: Download generated reports python3 scripts/data_agent_cli.py reports --session-id abc123xyz ``` **Benefits of Reuse**: - Avoid repeated data understanding phase - Preserve context history - Reduce API calls > See [ANALYSIS_MODE.md](ANALYSIS_MODE.md) for detailed sub-agent implementation specifications FILE:scripts/cli/__init__.py """CLI package for Data Agent unified command-line tool. Author: Tinker Created: 2026-03-04 """ from cli.formatters import ( _fmt_jupyter_cell, _fmt_task_finish, _fmt_insights, _fmt_table_summaries, _extract_json_objects, _format_data_event, _format_parsed_json, _fmt_plan_progress, _fmt_status_change, _fmt_output_conclusion, _fmt_recommended_questions, _SKIP_DATA_CATEGORIES, ) from cli.streaming import ( StreamState, _print_event, _stream_response, _finalize_stream, _is_user_confirmation_event, ) from cli.cmd_db import cmd_db, _build_data_source, _db_attach, _db_batch, _db_single from cli.cmd_ls import cmd_ls, _extract_list from cli.cmd_file import cmd_file, _print_generated_files from cli.cmd_attach import cmd_attach from cli.parser import build_parser, main __all__ = [ # Formatters "_fmt_jupyter_cell", "_fmt_task_finish", "_fmt_insights", "_fmt_table_summaries", "_extract_json_objects", "_format_data_event", "_format_parsed_json", "_fmt_plan_progress", "_fmt_status_change", "_fmt_output_conclusion", "_fmt_recommended_questions", "_SKIP_DATA_CATEGORIES", # Streaming "StreamState", "_print_event", "_stream_response", "_finalize_stream", "_is_user_confirmation_event", # Commands "cmd_db", "_build_data_source", "_db_attach", "_db_batch", "_db_single", "cmd_ls", "_extract_list", "cmd_file", "_print_generated_files", "cmd_attach", # Parser / entry "build_parser", "main", ] FILE:scripts/cli/cmd_attach.py """Attach to existing session subcommand (attach). Author: Tinker Created: 2026-03-04 """ import argparse import json import os import subprocess import sys from pathlib import Path from cli.streaming import _stream_response, StreamState, _print_event from cli.cmd_db import _db_attach from cli.log_handler import StructuredLogHandler # from cli.notify import push_notification from cli.worker_utils import is_worker_process, setup_async_worker, handle_worker_completion from cli.streaming_utils import run_worker_with_handler, execute_single_query from cli.worker_lock import check_worker_lock, write_worker_pid, acquire_worker_lock, release_worker_lock from data_agent import ( DataAgentConfig, DataAgentClient, SessionManager, MessageHandler, FileManager, SSEClient, ) def cmd_attach(args: argparse.Namespace) -> None: """Connect to an existing session for continuing analysis or confirming plan.""" session_id = args.session_id is_worker = os.environ.get("DATA_AGENT_ASYNC_WORKER") == "1" async_run = getattr(args, "async_run", True) # Initialize components config = DataAgentConfig.from_env() client = DataAgentClient(config) session_manager = SessionManager(client) message_handler = MessageHandler(client) sse_client = SSEClient(config) file_manager = FileManager(client) # Async mode is determined by the --async-run/--no-async-run flag, regardless of query presence if async_run and not is_worker: # PARENT PROCESS: spawn background worker print(f"Connecting to session: {session_id}") try: session = client.describe_session(session_id=session_id) except Exception as e: print(f"Error: Failed to connect to session: {e}", file=sys.stderr) sys.exit(1) # Manually create a temporary session object to pass to the common utility class TempSession: def __init__(self, sess_obj): self.session_id = sess_obj.session_id self.agent_id = sess_obj.agent_id temp_session = TempSession(session) # Use common async worker setup setup_async_worker(args, temp_session) elif is_worker: # WORKER PROCESS: run the attach operation in background - using common utility def attach_query_executor(message_handler, session, args): output_mode = getattr(args, "output", "summary") session_dir = Path(f"sessions/{session.session_id}") query = args.query if query: print(f"\n> User Query: {query}\n", flush=True) got_content, need_confirm, output_text = _stream_response( message_handler, session, query, output_mode=output_mode, output_dir=session_dir, process_log_handler=None, # In worker process, no process log handler needed for this path is_attach=True # Indicate this is an attach operation ) # Check if the query was for confirmation and if there are more steps to process is_confirmation_query = query.strip() in ["确认执行", "confirm", "execute", "同意后续所有SQL执行"] # If the query was a confirmation, there might be more steps to execute if is_confirmation_query: # After confirmation, the process has handled the confirmation # The subsequent steps will be tracked by the server and can be viewed # by attaching to the session again print(f"\n[Worker] Confirmation processed. Subsequent steps will continue in the background.", flush=True) # Determine final status final_status = "waiting_input" if need_confirm else "completed" with open(session_dir / "status.txt", "w", encoding="utf-8") as f: f.write(final_status) with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({ "status": final_status, "session_id": session.session_id, }, f, ensure_ascii=False, indent=2) return got_content, need_confirm else: # For attach without query, just monitor or replay from_start = getattr(args, "from_start", False) checkpoint = getattr(args, "checkpoint", None) # In worker process, we need to initialize components locally _config = DataAgentConfig.from_env() _client = DataAgentClient(_config) _sse_client = SSEClient(_config) _file_manager = FileManager(_client) _db_attach(_sse_client, _file_manager, session, from_start=from_start, checkpoint=checkpoint, output_mode=output_mode) return True, False # got_content=True, need_confirm=False # Since the attach worker has complex setup logic that differs significantly from the other commands, # we'll just call the run_worker_with_handler without a data_source run_worker_with_handler(args, query_execution_func=attach_query_executor) else: # SYNCHRONOUS MODE (--no-async-run) print(f"Connecting to session: {session_id}") print(f" Region: {config.region}") try: # First, try to get the session info to discover the actual agent_id # For some sessions, the API may not return the agent_id if an empty agent_id is provided # We need to try with empty agent_id first, and if it fails, we'll try the create_or_reuse approach session = None try: session = client.describe_session(agent_id="", session_id=session_id) # Special case: If we get CREATING status with no agent ID, this means the session doesn't actually exist # for the current user, but the API returns a default response. We should exit immediately. if not session.agent_id and session.status.value == "CREATING": print(f"Error: Session {session_id} does not exist or does not belong to the current user.", file=sys.stderr) sys.exit(1) # Only print session info if it passes the above check print(f"Initial session check - agent: '{session.agent_id}', status: {session.status.value}") except Exception as e: # This could be an authentication error or the session doesn't exist print(f"Error connecting to session: {e}", file=sys.stderr) print(f"This could be due to invalid credentials or the session doesn't exist.", file=sys.stderr) # If it's an authentication error, we should exit immediately from data_agent.exceptions import AuthenticationError if isinstance(e, AuthenticationError): print(f"Authentication failed. Please check your credentials in .env file.", file=sys.stderr) sys.exit(1) # For other errors, continue with fallback approaches print(f"Trying alternative methods to connect to session...", file=sys.stderr) # If session retrieval failed or agent_id is still empty, try using session manager if not session or not session.agent_id: print("Agent ID not found in session info, this may indicate the session doesn't exist or isn't ready yet.") # Check if the session exists but is still in CREATING state if session and session.status.value == "CREATING": print("Session exists but is still being created. Waiting for it to be ready...") # Since session is in CREATING state but no agent_id was returned, # we need to wait until the agent_id becomes available import time max_retries = 10 # Reduce wait time (10*30 = 5 minutes max) retry_count = 0 while retry_count < max_retries: time.sleep(30) # Wait 30 seconds between checks try: session = client.describe_session(agent_id="", session_id=session_id) print(f"Checking session again - agent: '{session.agent_id}', status: {session.status.value}") # If we now have an agent_id, break out of the loop if session.agent_id: print(f"Got agent ID: {session.agent_id}. Session is ready.") break # If status is no longer CREATING but we still don't have an agent_id, # the session may have failed to create properly or doesn't belong to the user if session.status.value != "CREATING": print(f"Session status changed to {session.status.value} but still no agent ID.") print(f"This may indicate the session does not belong to the current user.", file=sys.stderr) break except Exception as api_error: print(f"Check session retry failed: {api_error}", file=sys.stderr) from data_agent.exceptions import AuthenticationError if isinstance(api_error, AuthenticationError): print(f"Authentication failed. Please check your credentials in .env file.", file=sys.stderr) sys.exit(1) retry_count += 1 if not session.agent_id: print(f"Session is still not ready after waiting. Status: {session.status.value if session else 'unknown'}", file=sys.stderr) print(f"If status is not CREATING, this may mean the session does not belong to your account.", file=sys.stderr) if not session or not session.agent_id: # At this point, if we still don't have an agent_id, the session either: # 1. Does not exist # 2. Does not belong to the current user # 3. Is still being created but taking too long # 4. Has been cancelled or failed to create properly print(f"Error: Could not retrieve agent ID for session {session_id}.", file=sys.stderr) print(f"This may mean the session does not exist, does not belong to your account,", file=sys.stderr) print(f"is still being created (waited too long), or was cancelled/failed to create.", file=sys.stderr) sys.exit(1) rid = f", request_id: {session.request_id}" if session.request_id else "" print(f"Session connected: {session.session_id} (agent: {session.agent_id}, status: {session.status.value}{rid})") # If status is CREATING and a query is provided, we need to wait for it to be ready # But for IDLE state, we can proceed without waiting since it's a valid state if session.status.value == "CREATING" and getattr(args, "query", None): print("Waiting for session to be ready...") session = session_manager.create_or_reuse( session_id=session_id, agent_id=session.agent_id, wait_for_running=True, ) rid = f", request_id: {session.request_id}" if session.request_id else "" print(f"Session ready: {session.session_id} (status: {session.status.value}{rid})") except Exception as e: print(f"Error: Failed to connect to session: {e}", file=sys.stderr) sys.exit(1) output_mode = getattr(args, "output", "summary") session_dir = Path(f"sessions/{session.session_id}") # Create process log handler for attach command with StructuredLogHandler(session_dir, log_prefix="process") as log_handler: # Log initial connection info log_text = f"Connected to session: {session.session_id} (agent: {session.agent_id}, status: {session.status.value})" log_handler.write_both(log_text + "\n") if args.query: query = args.query print(f"\n> User Query: {query}\n") # Log the query query_log = f"Processing query: {query}\n" log_handler.write_both(query_log) try: got_content, need_confirm, _ = _stream_response( message_handler, session, query, output_mode=output_mode, output_dir=session_dir, process_log_handler=log_handler, # Pass the process log handler is_attach=True # Indicate this is an attach operation ) # Check if the query was for confirmation and if there are more steps to process is_confirmation_query = query.strip() in ["确认执行", "confirm", "execute", "同意后续所有SQL执行"] # Log the response completion response_log = f"Query processed, got_content: {got_content}, need_confirm: {need_confirm}, is_confirmation: {is_confirmation_query}\n" log_handler.write_both(response_log) if not got_content: print("(No response received, please retry)") elif need_confirm: print("\n⚠️ Agent has created an execution plan. User confirmation required.") print(f" To continue: python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'your input'") elif is_confirmation_query: # After confirmation, continue monitoring for additional steps print(f"\n[Sync mode] Post-confirmation: Monitoring for additional steps...") # Refresh session to get the latest status updated_session = client.describe_session(agent_id=session.agent_id, session_id=session.session_id) # If session is still running after confirmation, continue monitoring if updated_session.status.value in ["RUNNING", "WAIT_INPUT"]: print(f"Session status after confirmation: {updated_session.status.value}") # Continue to attach and monitor the session print(f"\nContinuing to monitor session progress...") # Use already imported _db_attach and initialized clients _db_attach(sse_client, file_manager, updated_session, from_start=False, checkpoint=None, output_mode=output_mode) else: try: updated_session = client.describe_session(agent_id=session.agent_id, session_id=session.session_id) if updated_session.status.value == "WAIT_INPUT": print("\n⚠️ Agent has created an execution plan and is waiting for confirmation.") print(" Use -q option to confirm the plan or provide feedback:") print(f" python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'confirm'") print(f" python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'modify the plan'") except Exception: pass except Exception as e: error_msg = f"Request failed: {e}" print(error_msg) log_handler.write_both(error_msg + "\n") else: from_start = getattr(args, "from_start", False) checkpoint = getattr(args, "checkpoint", None) # Log the attachment without query attach_log = f"Attaching to session without query (from_start: {from_start}, checkpoint: {checkpoint})\n" log_handler.write_both(attach_log) # Additional verification: if session agent_id is still empty, try to retrieve it directly if not session.agent_id: print(f"No agent ID in current session object, attempting to refresh session info...", file=sys.stderr) try: # Try to get fresh session data directly from the API refreshed_session = client.describe_session(agent_id="", session_id=session_id) if refreshed_session.agent_id: session = refreshed_session print(f"Successfully retrieved agent ID: {refreshed_session.agent_id}") else: print(f"Error: Could not retrieve agent ID for session {session_id}. Session may not exist.", file=sys.stderr) sys.exit(1) except Exception as e: print(f"Error retrieving session info: {e}", file=sys.stderr) sys.exit(1) # Final check before calling _db_attach if not session.agent_id: print(f"Error: Cannot attach to session {session_id} without a valid agent ID.", file=sys.stderr) print(f"The session might not exist or may be corrupted.", file=sys.stderr) sys.exit(1) _db_attach(sse_client, file_manager, session, from_start=from_start, checkpoint=checkpoint, output_mode=output_mode) FILE:scripts/cli/cmd_db.py """Database analysis subcommand (db). Author: Tinker Created: 2026-03-04 """ import argparse import sys import os import json import subprocess from pathlib import Path from typing import Optional # from cli.notify import push_notification from cli.worker_lock import check_worker_lock, write_worker_pid, acquire_worker_lock, release_worker_lock from cli.worker_utils import is_worker_process, setup_async_worker, handle_worker_completion from cli.streaming_utils import run_worker_with_handler, execute_single_query, execute_query_batch from cli.streaming import _stream_response, _finalize_stream, StreamState, _print_event from data_agent import ( DataAgentConfig, DataAgentClient, SessionManager, MessageHandler, FileManager, DataSource, SSEClient, ) def _build_data_source(args: argparse.Namespace) -> DataSource: """Build DataSource from command-line arguments.""" tables = [t.strip() for t in args.tables.split(",")] if args.tables else [] table_ids = [t.strip() for t in args.table_ids.split(",")] if args.table_ids else [] return DataSource( dms_instance_id=args.dms_instance_id, dms_database_id=args.dms_db_id, instance_name=args.instance_name, db_name=args.db_name, tables=tables, table_ids=table_ids, engine=args.engine, region_id=args.region, ) def _db_attach( sse_client: SSEClient, file_manager: FileManager, session, from_start: bool = False, checkpoint: Optional[int] = None, output_mode: str = "summary", ) -> None: """Attach to an existing session's SSE stream without sending a message. Streams all incoming events to stdout in real-time using the same formatting as _stream_response. Useful for: - Watching an ongoing analysis - Replaying the last round to review a plan before confirming After the stream ends, calls ListFileUpload to download any agent-generated report files to sessions//reports/. """ session_id = session.session_id session_dir = Path(f"sessions/{session_id}") if checkpoint is None: checkpoint = 0 if from_start else None label = "watching live stream" if from_start: label = "replaying from start" elif checkpoint is not None: label = f"resuming from checkpoint {checkpoint}" print(f"\nAttaching to session {session_id} ({label})...") print(f"Session status: {session.status.value}") # Special reminder for WAIT_INPUT status if session.status.value == "WAIT_INPUT": print("\n⚠️ Session is in WAIT_INPUT state.") print(" The agent has generated SQL and is waiting for your confirmation.") print() print(" To view the SQL and confirm:") print(f" python3 skill/data_agent_cli.py attach --session-id {session_id} --from-start") print() print(" To confirm and execute ONLY this SQL (DO NOT create a new session):") print(f" python3 skill/data_agent_cli.py attach --session-id {session_id} -q '确认执行当前SQL'") print() print(" To agree to execute all subsequent SQL automatically:") print(f" python3 skill/data_agent_cli.py attach --session-id {session_id} -q '同意后续所有SQL执行'") print() print(" To modify the query:") print(f" python3 skill/data_agent_cli.py attach --session-id {session_id} -q 'your new question'") print(f"Output directory: {session_dir.resolve()}") print("Press Ctrl+C to detach.") print("\n---\n") got_content = False last_checkpoint = checkpoint if checkpoint else 0 last_progress_time = 0 import time state = StreamState(output_mode=output_mode) state.output_dir = session_dir state.session_id = session_id state.is_attach = True state.session_status = getattr(session, 'status', None) if hasattr(state.session_status, 'value'): state.session_status = state.session_status.value try: # Initialize structured logging for attach mode from cli.streaming import init_structured_logging, close_structured_logging init_structured_logging(session_dir) for event in sse_client.stream_chat_content( agent_id=session.agent_id, session_id=session_id, checkpoint=checkpoint, ): if event.event_type == "SSE_FINISH": break # Track checkpoint progress (only in detail/raw mode) if output_mode != "summary" and event.checkpoint is not None and event.checkpoint > last_checkpoint: current_time = time.time() # Show checkpoint progress every 30 seconds or every 50 checkpoints if current_time - last_progress_time >= 30 or event.checkpoint - last_checkpoint >= 50: print(f" [checkpoint: {event.checkpoint}]", flush=True) last_progress_time = current_time last_checkpoint = event.checkpoint c, _ = _print_event(event, output_mode, state=state) if c: got_content = True except KeyboardInterrupt: print("\n\nDetached.") return except Exception as e: # Extract request_id from HTTP error response if available request_id = "" resp = getattr(e, "response", None) if resp is not None: request_id = resp.headers.get("x-acs-request-id", "") rid_str = f" (Request-Id: {request_id})" if request_id else "" print(f"\nError: {e}{rid_str}", file=sys.stderr) return finally: # Close structured logging close_structured_logging() _finalize_stream(state) if state.got_content: got_content = True if got_content: print() # trailing newline else: print("(No new events -- session may be waiting for your input)") # -- Download agent-generated files via ListFileUpload -- # Server needs a few seconds to finalize report files after session completes print("\nFetching agent-generated files (waiting for server to finalize)...") time.sleep(5) report_dir = session_dir / "reports" total_reports = 0 for category in ("WebReport", "TextReport", "DefaultArtifact"): try: files = file_manager.list_files(session_id, file_category=category) except Exception as e: print(f" Warning: could not list {category}: {e}", file=sys.stderr) continue if not files: continue report_dir.mkdir(parents=True, exist_ok=True) for rf in files: if not rf.download_url: print(f" Skipping {rf.filename} ({category}): no download URL") continue save_path = report_dir / (rf.filename or f"{rf.file_id}.bin") try: file_manager.download_from_url(rf.download_url, str(save_path)) print(f" [{category}] saved \u2192 {save_path.resolve()}") total_reports += 1 except Exception as e: print(f" Failed to download {rf.filename} ({category}): {e}", file=sys.stderr) if total_reports == 0: print(" No report files found for this session.") print("\n---\n") print(f'> \U0001f4a1 To continue conversation:') print(f'> python3 skill/data_agent_cli.py attach --session-id {session_id} -q "your message"') def _db_batch( message_handler: MessageHandler, session, data_source: DataSource, queries: list, output_mode: str = "summary", output_dir: Optional[Path] = None, ) -> tuple[bool, bool, str]: """Execute batch preset queries.""" got_content, need_confirm = False, False full_text = "" for query in queries: print(f"\n{'=' * 60}") print(f"Query: {query}") print("=" * 60) c, nc, t = _stream_response(message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir) if c: got_content = True if t: full_text += f"\n### Query: {query}\n" + t + "\n" if nc: need_confirm = True break return got_content, need_confirm, full_text def _db_single( message_handler: MessageHandler, session, data_source: DataSource, query: str, output_mode: str = "summary", output_dir: Optional[Path] = None, ) -> tuple[bool, bool, str]: """Execute a single query with streaming output.""" print(f"\nAnalyzing...\n") got_content, need_confirm, full_text = _stream_response( message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir ) if not got_content: print("(No response received, please retry)") elif need_confirm: print("\n\u26a0\ufe0f \u9700\u8981\u7528\u6237\u786e\u8ba4\uff0c\u7a0b\u5e8f\u5c06\u9000\u51fa\u3002\u8bf7\u5b8c\u6210\u786e\u8ba4\u540e\u4f7f\u7528\u4f1a\u8bddID\u7ee7\u7eed\u5bf9\u8bdd\u3002") return got_content, need_confirm, full_text def cmd_db(args: argparse.Namespace) -> None: """Handle db subcommand.""" # Validate required database parameters missing = [] for attr, name in [ ("dms_instance_id", "--dms-instance-id"), ("dms_db_id", "--dms-db-id"), ("instance_name", "--instance-name"), ("db_name", "--db-name"), ("tables", "--tables"), ]: if not getattr(args, attr, None): missing.append(name) if missing: print(f"Error: Missing required parameters: {', '.join(missing)}", file=sys.stderr) sys.exit(1) # Initialize components config = DataAgentConfig.from_env() client = DataAgentClient(config) session_manager = SessionManager(client) message_handler = MessageHandler(client) # Create new session session_mode = args.session_mode.upper() data_source = _build_data_source(args) is_worker = is_worker_process() if getattr(args, "async_run", False) and not is_worker: # PARENT PROCESS LOGIC enable_search = getattr(args, 'enable_search', False) print(f"Creating session for async execution...") session = session_manager.create_or_reuse(mode=session_mode, database_id=str(args.dms_db_id), enable_search=enable_search) # Use common async worker setup setup_async_worker(args, session) elif is_worker: # WORKER PROCESS LOGIC - using common utility def db_query_executor(message_handler, session, args): output_mode = getattr(args, "output", "summary") data_source = _build_data_source(args) if args.query: _, need_confirm, _ = _db_single(message_handler, session, data_source, args.query, output_mode=output_mode, output_dir=Path(f"sessions/{session.session_id}")) else: if session_mode == "ANALYSIS": default_queries = [ f"Analyze the overall data structure and table relationships of {data_source.db_name} database", "Identify key metrics and distribution characteristics in the data", ] else: default_queries = [ f"What tables exist in {data_source.db_name} database?", "Who has the highest sales?", ] _, need_confirm, _ = _db_batch(message_handler, session, data_source, default_queries, output_mode=output_mode, output_dir=Path(f"sessions/{session.session_id}")) return True, need_confirm # got_content=True, need_confirm run_worker_with_handler(args, data_source=_build_data_source(args), query_execution_func=db_query_executor) # NORMAL SYNCHRONOUS LOGIC mode_desc = { "ASK_DATA": "ASK_DATA mode (SQL query + natural language response)", "ANALYSIS": "ANALYSIS mode (deep analysis + report generation)", "INSIGHT": "INSIGHT mode", }.get(session_mode, session_mode) print(f"Creating session: {mode_desc}...") print(f" Region: {config.region}") enable_search = getattr(args, 'enable_search', False) session = session_manager.create_or_reuse(mode=session_mode, database_id=str(args.dms_db_id), enable_search=enable_search) print(f"Session ready: {session.session_id}") print(f"\n💡 Tip: To continue this session later, use: python3 scripts/data_agent_cli.py attach --session-id {session.session_id}") # Execute query output_mode = getattr(args, "output", "summary") session_dir = Path(f"sessions/{session.session_id}") session_dir.mkdir(parents=True, exist_ok=True) # Initialize structured logging for sync mode from cli.streaming import init_structured_logging, close_structured_logging init_structured_logging(session_dir) try: if args.query: _, _, output_text = _db_single(message_handler, session, data_source, args.query, output_mode=output_mode, output_dir=session_dir) else: # Default batch preset queries if session_mode == "ANALYSIS": default_queries = [ f"Analyze the overall data structure and table relationships of {data_source.db_name} database", "Identify key metrics and distribution characteristics in the data", ] else: default_queries = [ f"What tables exist in {data_source.db_name} database?", "Who has the highest sales?", ] print(f"\nNo query specified, running preset queries ({len(default_queries)} total)...") _, _, output_text = _db_batch(message_handler, session, data_source, default_queries, output_mode=output_mode, output_dir=session_dir) finally: close_structured_logging() # Write result status # if output_text: # with open(session_dir / "output.md", "w", encoding="utf-8") as f: # f.write(output_text) # with open(session_dir / "result.json", "w", encoding="utf-8") as f: # json.dump({"status": "completed", "output_file": "output.md"}, f) with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({"status": "completed"}, f) FILE:scripts/cli/cmd_dms.py """DMS tools subcommand (dms). Author: Tinker Created: 2026-03-05 """ import argparse import sys from data_agent import DataAgentConfig, DmsMcpTools def cmd_dms(args: argparse.Namespace) -> None: """Handle dms subcommand for DMS tools.""" config = DataAgentConfig.from_env() mcp_tools = DmsMcpTools(config) tool = getattr(args, "tool", None) if tool == "list-instances": _list_instances(mcp_tools, args) elif tool == "search-database": _search_database(mcp_tools, args) elif tool == "list-tables": _list_tables(mcp_tools, args) else: print("Error: Unknown tool. Use 'list-instances', 'search-database', or 'list-tables'.", file=sys.stderr) sys.exit(1) def _list_instances(mcp_tools: DmsMcpTools, args: argparse.Namespace) -> None: """List DMS instances using MCP tool.""" search_key = getattr(args, "search", None) db_type = getattr(args, "db_type", None) env_type = getattr(args, "env_type", None) page_number = getattr(args, "page_number", 1) page_size = getattr(args, "page_size", 50) print(f"Region: {mcp_tools._config.region}") print("Fetching instances from DMS...") print("-" * 60) try: result = mcp_tools.list_instances( search_key=search_key, db_type=db_type, env_type=env_type, page_number=page_number, page_size=page_size, ) if not result.items: print("No instances found.") return print(f"Page {result.page_number}/{result.total_pages} (Total: {result.total_count} instances)\n") for inst in result.items: print(f" {inst.instance_alias} [{inst.instance_type}]") print(f" Instance ID : {inst.instance_id}") print(f" Host:Port : {inst.host}:{inst.port}") print(f" State : {inst.state}") print(f" Env Type : {inst.env_type}") print(f" Source : {inst.instance_source}") if inst.instance_resource_id: print(f" Resource ID : {inst.instance_resource_id}") print() # Show pagination hint print("-" * 60) if result.has_next: print(f"📄 Next page: --page-number {result.page_number + 1}") else: print("📄 This is the last page.") except Exception as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) def _search_database(mcp_tools: DmsMcpTools, args: argparse.Namespace) -> None: """Search databases by schema name.""" search_key = getattr(args, "search_key", None) page_number = getattr(args, "page_number", 1) page_size = getattr(args, "page_size", 200) if not search_key: print("Error: --search-key is required for search-database tool.", file=sys.stderr) sys.exit(1) print(f"Region: {mcp_tools._config.region}") print(f"Search Key: {search_key}") print("Searching databases in DMS...") print("-" * 60) try: databases = mcp_tools.search_database( search_key=search_key, page_number=page_number, page_size=page_size, ) if not databases: print("No databases found.") print("\nTroubleshooting tips:") print(" 1. The database may not be registered in DMS") print(" 2. The search_key is case-sensitive, try different cases") print(" 3. Try using 'dms list-instances' to see available instances first") return print(f"Found {len(databases)} database(s):\n") for db in databases: print(f" {db.schema_name} [{db.db_type}]") print(f" Database ID : {db.database_id}") print(f" Host:Port : {db.host}:{db.port}") print(f" Instance : {db.instance_alias} (ID: {db.instance_id})") print(f" Env Type : {db.env_type}") print() except Exception as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) def _list_tables(mcp_tools: DmsMcpTools, args: argparse.Namespace) -> None: """List tables in a database.""" database_id = getattr(args, "database_id", None) search_name = getattr(args, "search_name", None) page_number = getattr(args, "page_number", 1) page_size = getattr(args, "page_size", 200) if not database_id: print("Error: --database-id is required for list-tables tool.", file=sys.stderr) sys.exit(1) print(f"Region: {mcp_tools._config.region}") print(f"Database ID: {database_id}") if search_name: print(f"Search Name: {search_name}") print("Fetching tables from DMS...") print("-" * 60) try: tables = mcp_tools.list_tables( database_id=database_id, search_name=search_name, page_number=page_number, page_size=page_size, ) if not tables: print("No tables found.") return print(f"Found {len(tables)} table(s):\n") for table in tables: print(f" {table.table_name}") print(f" Table ID : {table.table_id}") print(f" Table GUID : {table.table_guid}") print(f" Schema : {table.schema_name}") if table.engine: print(f" Engine : {table.engine}") if table.table_comment: print(f" Comment : {table.table_comment}") print() except Exception as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) FILE:scripts/cli/cmd_file.py """File analysis subcommand (file). Author: Tinker Created: 2026-03-04 """ import argparse import sys import os import json import subprocess from pathlib import Path # from cli.notify import push_notification from cli.worker_utils import is_worker_process, setup_async_worker, handle_worker_completion from cli.streaming_utils import run_worker_with_handler, execute_single_query, execute_query_batch from cli.worker_lock import check_worker_lock, write_worker_pid, acquire_worker_lock, release_worker_lock from cli.dual_logger import run_with_dual_logging from cli.streaming import _stream_response from data_agent import ( DataAgentConfig, DataAgentClient, SessionManager, MessageHandler, FileManager, DataSource, ) def cmd_file(args: argparse.Namespace) -> None: """Handle file subcommand.""" file_id = getattr(args, 'file_id', None) file_path = args.file_path # Validate arguments if not file_id and not file_path: print("Error: Either FILE path or --file-id must be provided", file=sys.stderr) sys.exit(1) if file_id and file_path: print("Error: Cannot specify both FILE path and --file-id", file=sys.stderr) sys.exit(1) # Initialize components config = DataAgentConfig.from_env() client = DataAgentClient(config) session_manager = SessionManager(client) message_handler = MessageHandler(client) file_manager = FileManager(client) is_worker = is_worker_process() session_mode = args.session_mode.upper() if getattr(args, "async_run", False) and not is_worker: # PARENT PROCESS LOGIC enable_search = getattr(args, 'enable_search', False) print(f"Creating session for async file analysis...") if file_id: # Using Data Center file ID print(f"Using Data Center file ID: {file_id}") # Create data source with file ID for query execution file_data_source = DataSource( data_source_type="FILE", file_id=file_id, region_id=config.region ) # Create regular session without binding to specific data source initially session = session_manager.create_or_reuse(mode=session_mode, enable_search=enable_search, file_id=file_id) else: # Local file workflow: upload first if not Path(file_path).exists(): print(f"Error: File not found: {file_path}", file=sys.stderr) sys.exit(1) # Validate file type if not file_manager.is_supported_file(file_path): supported = ".csv, .xlsx, .xls" print(f"Error: Unsupported file type. Supported formats: {supported}", file=sys.stderr) sys.exit(1) # Upload file and get file info print(f"Uploading file: {file_path}") try: file_info = file_manager.upload_file(file_path) print(f"File uploaded successfully!") print(f" File ID : {file_info.file_id}") print(f" Filename: {file_info.filename}") print(f" Size : {file_info.size} bytes") # Create data source with uploaded file ID file_data_source = DataSource( data_source_type="FILE", file_id=file_info.file_id, ) except Exception as e: print(f"Upload failed: {e}", file=sys.stderr) sys.exit(1) # Create session session = session_manager.create_or_reuse(mode=session_mode, enable_search=enable_search, file_id=file_info.file_id) # Store file_data_source as a dict for async worker to use (avoid serialization issues) args.file_data_source = file_data_source.to_api_dict() # Use common async worker setup setup_async_worker(args, session) print(f"\n✅ Async task started. Session ID: {session.session_id}") print(f"Check progress at: sessions/{session.session_id}/progress.log") sys.exit(0) elif is_worker: # WORKER PROCESS LOGIC - using common utility def file_query_executor(message_handler, session, args): if hasattr(args, 'file_data_source'): # If file_data_source is a dict (async mode), reconstruct DataSource if isinstance(args.file_data_source, dict): # Create a new DataSource from the dict representation # For file analysis, we need to use data_source_type="FILE" and the correct file_id # The DataSource.to_api_dict() method expects data_source_type="FILE" to trigger FileId inclusion file_data_source = DataSource( data_source_type="FILE", # Use "FILE" type for file analysis file_id=args.file_data_source.get('FileId', ''), region_id=args.file_data_source.get('RegionId', config.region) ) else: # If file_data_source is already a DataSource object (sync mode) file_data_source = args.file_data_source else: # Try to load from input.json if not available in args # In worker process, we can determine the session dir and load input.json import json from pathlib import Path session_dir = Path(f"sessions/{session.session_id}") input_file = session_dir / "input.json" if input_file.exists(): try: with open(input_file, 'r', encoding='utf-8') as f: saved_args = json.load(f) if 'file_data_source' in saved_args: file_data_source_dict = saved_args['file_data_source'] # Create DataSource for file analysis using data_source_type="FILE" file_data_source = DataSource( data_source_type="FILE", # Use "FILE" type for file analysis file_id=file_data_source_dict.get('FileId', ''), region_id=file_data_source_dict.get('RegionId', config.region) ) except Exception as e: print(f"Warning: Could not load file_data_source from input.json: {e}", file=sys.stderr) if file_data_source is None: # Fallback: reconstruct based on available info if hasattr(args, 'file_id') and args.file_id: file_data_source = DataSource( data_source_type="FILE", file_id=args.file_id, region_id=getattr(args, 'region_id', 'cn-hangzhou') ) else: print("Error: No file data source available in worker", file=sys.stderr) return False, False output_mode = getattr(args, "output", "summary") session_dir = Path(f"sessions/{session.session_id}") if args.query: queries = [args.query] else: queries = [ "请分析上传文件的数据结构", "数据的关键统计指标和分布情况是什么？", "数据中是否存在异常值或离群点？", ] print(f"\n{'=' * 60}") print("File Analysis") print("=" * 60) output_text = "" for query in queries: print(f"\nQuery: {query}") print("-" * 50) got_content, need_confirm, t = _stream_response( message_handler, session, query, data_source=file_data_source, output_mode=output_mode, output_dir=session_dir, ) if t: output_text += f"\n### Query: {query}\n" + t + "\n" if need_confirm: return got_content, need_confirm # got_content, need_confirm if args.list_generated_files: _print_generated_files(file_manager, session.session_id) return True, False # got_content=True, need_confirm=False # Pass args as-is since we've attached file_data_source to it run_worker_with_handler(args, query_execution_func=file_query_executor) # NORMAL SYNCHRONOUS LOGIC if file_id: # Using Data Center file ID print(f"Using Data Center file ID: {file_id}") # Create data source with file ID file_data_source = DataSource( data_source_type="FILE", file_id=file_id, region_id=config.region ) else: # Local file workflow: validate and upload if not Path(file_path).exists(): print(f"Error: File not found: {file_path}", file=sys.stderr) sys.exit(1) # Validate file type if not file_manager.is_supported_file(file_path): supported = ".csv, .xlsx, .xls" print(f"Error: Unsupported file type. Supported formats: {supported}", file=sys.stderr) sys.exit(1) # Upload file print(f"Uploading file: {file_path}") try: file_info = file_manager.upload_file(file_path) except Exception as e: print(f"Upload failed: {e}", file=sys.stderr) sys.exit(1) print(f"File uploaded successfully!") print(f" File ID : {file_info.file_id}") print(f" Filename: {file_info.filename}") print(f" Size : {file_info.size} bytes") # Create data source for uploaded file file_data_source = DataSource( data_source_type="FILE", file_id=file_info.file_id, ) # Create session session_mode = args.session_mode.upper() mode_desc = { "ASK_DATA": "ASK_DATA mode", "ANALYSIS": "ANALYSIS mode (recommended for file analysis)", "INSIGHT": "INSIGHT mode", }.get(session_mode, session_mode) print(f"\nCreating session: {mode_desc}...") print(f" Region: {config.region}") enable_search = getattr(args, 'enable_search', False) # For file analysis, pass file_id to create_or_reuse so the session is bound to the file session = session_manager.create_or_reuse(mode=session_mode, enable_search=enable_search, file_id=file_data_source.file_id if 'file_data_source' in locals() else None) print(f"Session ready: {session.session_id}") print(f"\n💡 Tip: To continue this session later, use: python3 data_agent_cli.py attach --session-id {session.session_id}") # Get output mode output_mode = getattr(args, "output", "summary") session_dir = Path(f"sessions/{session.session_id}") session_dir.mkdir(parents=True, exist_ok=True) # Initialize structured logging for sync mode from cli.streaming import init_structured_logging, close_structured_logging init_structured_logging(session_dir) try: # Determine queries to execute if args.query: queries = [args.query] else: # Default preset analysis questions queries = [ "请分析上传文件的数据结构", "数据的关键统计指标和分布情况是什么？", "数据中是否存在异常值或离群点？", ] # Execute queries with the file data source print(f"\n{'=' * 60}") print("File Analysis") print("=" * 60) output_text = "" for query in queries: print(f"\nQuery: {query}") print("-" * 50) try: got_content, need_confirm, t = _stream_response( message_handler, session, query, data_source=file_data_source, output_mode=output_mode, output_dir=session_dir, ) if t: output_text += f"\n### Query: {query}\n" + t + "\n" if not got_content: print("(No response received, please retry)") elif need_confirm: print("\n⚠️ Agent has created an execution plan. User confirmation required.") print(f" To continue: python3 scripts/data_agent_cli.py attach --session-id {session.session_id} -q 'User Input' ") if output_text: with open(session_dir / "output.md", "w", encoding="utf-8") as f: f.write(output_text) with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({"status": "waiting_input", "output_file": "output.md"}, f) return except Exception as e: print(f"Request failed: {e}") # List generated files if requested if args.list_generated_files: _print_generated_files(file_manager, session.session_id) finally: # Close structured logging close_structured_logging() # Write result status with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({"status": "completed"}, f) def _print_generated_files(file_manager: FileManager, session_id: str) -> None: """Print list of files generated by the Agent.""" print(f"\n{'=' * 60}") print("Generated Files") print("=" * 60) try: generated = file_manager.list_files(session_id) if generated: for f in generated: print(f" - {f.filename} ({f.file_type}, {f.size} bytes)") if f.download_url: print(f" Download: {f.download_url}") else: print(" No generated files.") except Exception as e: print(f" Failed to get file list: {e}") FILE:scripts/cli/cmd_import.py """Import subcommand - Add DMS database tables to Data Agent Data Center. Author: Tinker Created: 2026-03-05 """ import argparse import sys from data_agent import DataAgentConfig, DataAgentClient def cmd_import(args: argparse.Namespace) -> None: """Handle import subcommand for adding DMS tables to Data Center.""" # Validate required parameters missing = [] for attr, name in [ ("dms_instance_id", "--dms-instance-id"), ("dms_db_id", "--dms-db-id"), ("instance_name", "--instance-name"), ("db_name", "--db-name"), ("tables", "--tables"), ]: if not getattr(args, attr, None): missing.append(name) if missing: print(f"Error: Missing required parameters: {', '.join(missing)}", file=sys.stderr) sys.exit(1) # Initialize client config = DataAgentConfig.from_env() client = DataAgentClient(config) # Parse tables table_list = [t.strip() for t in args.tables.split(",")] print("Adding database tables to Data Center...") print(f" Region: {config.region}") print(f" Instance: {args.instance_name}") print(f" Database: {args.db_name}") print(f" DMS Instance ID: {args.dms_instance_id}") print(f" DMS DB ID: {args.dms_db_id}") print(f" DB Type: {getattr(args, 'engine', 'mysql')}") print(f" Tables: {', '.join(table_list)}") print("-" * 60) try: result = client.add_data_center_table( instance_name=args.instance_name, database_name=args.db_name, dms_instance_id=args.dms_instance_id, dms_db_id=args.dms_db_id, table_name_list=table_list, db_type=getattr(args, "engine", "mysql"), region_id=args.region, ) print("✓ Successfully added to Data Center") print(f" Response: {result.get('Data', result)}") except Exception as e: print(f"✗ Failed to add to Data Center: {e}", file=sys.stderr) sys.exit(1) FILE:scripts/cli/cmd_ls.py """List databases and tables (ls subcommand). Author: Tinker Created: 2026-03-04 """ import argparse import sys from data_agent import DataAgentConfig, DataAgentClient def _extract_list(resp: dict) -> list: """Extract a list from an API response regardless of nesting style. Handles ``{Data: [...]}`` and ``{Data: {List/DataList/Content: [...]}}``. Also handles lowercase ``{data: {...}}`` format from some API responses. """ # Try uppercase "Data" first, then lowercase "data" data = resp.get("Data") or resp.get("data", []) if isinstance(data, list): return data if isinstance(data, dict): return data.get("List") or data.get("DataList") or data.get("Content") or [] return [] def _get_field(obj: dict, *names: str, default=""): """Get a field value trying multiple possible key names (case-insensitive). Args: obj: The dictionary to search *names: Possible field names to try (e.g., "DatabaseName", "databaseName") default: Default value if none found """ for name in names: if name in obj: return obj[name] return default def cmd_ls(args: argparse.Namespace) -> None: """List DMS databases and (optionally) their tables.""" config = DataAgentConfig.from_env() client = DataAgentClient(config) search = getattr(args, "search", None) db_id = getattr(args, "db_id", None) sep = "-" * 60 print(f"Region: {config.region}") # -- list databases -- if db_id is None: print("Fetching databases...") try: resp = client.list_databases(search_key=search) except Exception as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) items = _extract_list(resp) if not items: print("No databases found.") return # Separate real DB connections from file-based data sources # Handle both PascalCase and camelCase keys from API response def _get_import_type(d): return d.get("ImportType") or d.get("importType", "") real_dbs = [d for d in items if _get_import_type(d) in ("RDS", "DMS")] file_dbs = [d for d in items if _get_import_type(d) == "FILE"] # -- Real databases (RDS / DMS) -- print(f"\n{'=' * 60}") print(f" Database Connections ({len(real_dbs)}) [ImportType: RDS/DMS]") print(f"{'=' * 60}") if real_dbs: for db in real_dbs: db_name = _get_field(db, "DatabaseName", "databaseName") db_type = _get_field(db, "DbType", "dbType", default="").lower() import_type = _get_field(db, "ImportType", "importType") dms_db_id = _get_field(db, "DmsDbId", "dmsDbId") dms_instance_id = _get_field(db, "DmsInstanceId", "dmsInstanceId") instance_name = _get_field(db, "InstanceName", "instanceName") db_desc = _get_field(db, "DatabaseDesc", "databaseDesc") agent_db_id = _get_field(db, "DbId", "dbId") print(f"\n {db_name} [{db_type}] ({import_type})") if db_desc: print(f" Desc : {db_desc}") print(f" AgentDbId : {agent_db_id}") print(f" DmsDbId : {dms_db_id}") print(f" DmsInstanceId : {dms_instance_id}") print(f" InstanceName : {instance_name}") else: print(" (none)") # -- File-based data sources -- print(f"\n{'=' * 60}") print(f" File Data Sources ({len(file_dbs)}) [ImportType: FILE]") print(f"{'=' * 60}") if file_dbs: for db in file_dbs: db_name = _get_field(db, "DatabaseName", "databaseName") db_type = _get_field(db, "DbType", "dbType", default="").lower() agent_db_id = _get_field(db, "DbId", "dbId") db_desc = _get_field(db, "DatabaseDesc", "databaseDesc") internal = _get_field(db, "IsInternal", "isInternal", default="N") label = "[sample]" if internal == "Y" else "" print(f" {db_name:<45} [{db_type}] {label} DbId={agent_db_id}") else: print(" (none)") print() return # -- list tables for a specific db_id -- print(f"Fetching tables for DbId={db_id}...") # Fetch all databases first to get InstanceName + DatabaseName (required by API) db_meta: dict = {} try: db_resp = client.list_databases() all_dbs = _extract_list(db_resp) for db in all_dbs: if str(db.get("DbId", "")) == str(db_id): db_meta = db break except Exception: pass if not db_meta: print(f"Error: DbId '{db_id}' not found in database list.", file=sys.stderr) sys.exit(1) inst_name = db_meta.get("InstanceName", "") db_name_q = db_meta.get("DatabaseName", "") try: resp = client.list_tables(inst_name, db_name_q) except Exception as e: print(f"Error: {e}", file=sys.stderr) sys.exit(1) items = _extract_list(resp) if not items: print("No tables found.") return table_names = [t.get("TableName", t.get("Name", "")) for t in items] table_ids = [t.get("TableId", t.get("Id", "")) for t in items] db_name = db_meta.get("DatabaseName", "") db_type = db_meta.get("DbType", "").lower() dms_instance_id = db_meta.get("DmsInstanceId", "") dms_db_id = db_meta.get("DmsDbId", "") instance_name = db_meta.get("InstanceName", "") print(f"\n{'=' * 60}") print(f" Database : {db_name} [{db_type}]") print(f" AgentDbId : {db_id}") print(f" DmsDbId : {dms_db_id}") print(f" Instance : {instance_name} (DmsInstanceId={dms_instance_id})") print(f" Tables : {len(items)}") print(f"{'=' * 60}") for name, tid in zip(table_names, table_ids): print(f" {name:<30} {tid}") print() # Print ready-to-use CLI command tables_arg = ",".join(table_names) table_ids_arg = ",".join(table_ids) print(sep) print(" Ready-to-use db command:") print(sep) print(f" python3 data_agent_cli.py db \\") print(f" --dms-instance-id {dms_instance_id} \\") print(f" --dms-db-id {dms_db_id} \\") print(f" --instance-name {instance_name} \\") print(f" --db-name {db_name} \\") print(f" --engine {db_type} \\") print(f" --tables {tables_arg} \\") print(f" --session-mode ASK_DATA \\") print(f" -q \"your question here\"") print(sep) FILE:scripts/cli/cmd_reports.py """Reports subcommand to list and download generated files. Author: Tinker Created: 2026-03-11 """ import argparse import sys from pathlib import Path import time from data_agent import ( DataAgentConfig, DataAgentClient, FileManager, ) def cmd_reports(args: argparse.Namespace) -> None: """Handle reports subcommand.""" session_id = args.session_id session_dir = Path(f"sessions/{session_id}") report_dir = session_dir / "reports" print(f"Fetching generated files for session {session_id}...") # Initialize components config = DataAgentConfig.from_env() client = DataAgentClient(config) file_manager = FileManager(client) total_reports = 0 categories = ("WebReport", "TextReport", "DefaultArtifact") found_files = [] for category in categories: try: files = file_manager.list_files(session_id, file_category=category) if files: found_files.extend([(category, rf) for rf in files]) except Exception as e: print(f" Warning: could not list {category}: {e}", file=sys.stderr) continue if not found_files: print("No report files found for this session.") print("Note: Reports are usually generated in ANALYSIS or INSIGHT modes.") return report_dir.mkdir(parents=True, exist_ok=True) print(f"Downloading files to {report_dir.resolve()}...\n") for category, rf in found_files: if not rf.download_url: print(f" [{category}] {rf.filename or rf.file_id}: No download URL available") continue save_path = report_dir / (rf.filename or f"{rf.file_id}.bin") try: print(f" Downloading [{category}] {rf.filename or rf.file_id}...") file_manager.download_from_url(rf.download_url, str(save_path)) print(f" ✅ Saved to {save_path.resolve()}") total_reports += 1 except Exception as e: print(f" ❌ Failed to download {rf.filename or rf.file_id} ({category}): {e}", file=sys.stderr) if total_reports > 0: print(f"\nSuccessfully downloaded {total_reports} files.") FILE:scripts/cli/dual_logger.py """Simple dual logging for Data Agent sessions. This module provides functionality to create both progress.log and progress.jsonl files. """ import subprocess from pathlib import Path from datetime import datetime import json from typing import Optional def run_with_dual_logging(cmd, session_dir: Path, env=None): """Run a subprocess and create both progress.log and progress.jsonl files.""" session_dir.mkdir(parents=True, exist_ok=True) # Run the subprocess with output redirected to a temporary file temp_output_path = session_dir / "temp_output.txt" with open(temp_output_path, 'w', encoding='utf-8') as temp_file: proc = subprocess.Popen( cmd, stdout=temp_file, stderr=subprocess.STDOUT, start_new_session=True, env=env ) proc.wait() # After subprocess completes, process the output to create both log files process_output_to_dual_logs(temp_output_path, session_dir) # Clean up temporary file temp_output_path.unlink(missing_ok=True) return proc def process_output_to_dual_logs(temp_output_path: Path, session_dir: Path): """Process the temporary output file and create log format only (progress.jsonl disabled).""" progress_log_path = session_dir / "progress.log" # progress_jsonl_path = session_dir / "progress.jsonl" # Disabled per user request with open(temp_output_path, 'r', encoding='utf-8') as temp_file: content = temp_file.read() # Write to progress.log as-is with open(progress_log_path, 'w', encoding='utf-8') as log_file: log_file.write(content) # JSONL logging disabled - progress.jsonl creation is skipped # Process content for progress.jsonl (skipped) FILE:scripts/cli/formatters.py """SSE data event formatting helpers for CLI display. Author: Tinker Created: 2026-03-04 """ import base64 import json import mimetypes import re from pathlib import Path from typing import Any, Optional # Categories to skip (mirrors SSEClient._SKIP_DATA_CATEGORIES) _SKIP_DATA_CATEGORIES = {"tool_call_choices", "metric_agent_config"} def _fmt_jupyter_cell(outer: dict) -> Optional[str]: """Format a jupyter_cell data event into human-readable text. Decision logic: - Parse nb_file_outputs for real content - Skip cells whose only outputs are dms/executing placeholders - SQL cells: show query + result table - Code cells: show stdout only (not the code itself) """ result_raw = outer.get("result", "") try: inner = json.loads(result_raw) if isinstance(result_raw, str) else result_raw except (json.JSONDecodeError, TypeError): return None title = inner.get("title", "") content_type = inner.get("content_type", "") cell_content = inner.get("content", "").strip() # Collect real output blocks first; skip this cell entirely if empty output_lines = [] for output in inner.get("nb_file_outputs", []): otype = output.get("output_type", "") metadata = output.get("metadata", {}) # Skip dms/executing placeholders ("正在执行中...") if otype == "display_data" and metadata.get("content_type") == "dms/executing": continue if otype == "display_data": app_json = ( output.get("data", {}) .get("application/json", {}) .get("data", {}) ) columns = app_json.get("columns", []) rows = app_json.get("result", []) if columns and rows: headers = [c.get("title", c.get("field", "?")) for c in columns] output_lines.append("Result:") output_lines.append(" " + " | ".join(headers)) output_lines.append(" " + "-" * max(len(" | ".join(headers)), 30)) for row in rows: vals = [str(row.get(c["field"], "")) for c in columns] output_lines.append(" " + " | ".join(vals)) elif otype == "stream": text = output.get("text", "").strip() if text: output_lines.append(f"Output:\n{text}") # Nothing real to show -> suppress this cell if not output_lines: return None lines = [] sep = "-" * 50 if title: lines.append(f"\n{sep}") lines.append(f" {title}") lines.append(sep) # SQL cells: show the query if content_type == "sql" and cell_content: lines.append(f"SQL:\n{cell_content}") # Code cells: suppress code text, show output only lines.extend(output_lines) return "\n".join(lines) def _fmt_task_finish(payload: dict) -> Optional[str]: """Format a task_finish data event for structured CLI display (ASK_DATA mode). Handles common field layouts returned by the Data Agent: - status / taskStatus / state - conclusion / summaryResult / summary / answer (natural language result) - sql / sqlList (executed SQL) - result / resultData / data (table with columns+rows) """ lines: list[str] = [] lines.append(f"\n### Query Result\n") # Status line status = ( payload.get("status") or payload.get("taskStatus") or payload.get("state") or payload.get("taskState") ) if status: lines.append(f" Status : {status}") # Natural language conclusion conclusion = ( payload.get("conclusion") or payload.get("summaryResult") or payload.get("summary") or payload.get("answer") or payload.get("content") or "" ) if conclusion and isinstance(conclusion, str) and conclusion.strip(): lines.append(f"\n{conclusion.strip()}") # SQL sql = payload.get("sql") if not sql: sql_list = payload.get("sqlList") or [] sql = sql_list[0] if sql_list else None if sql and isinstance(sql, str) and sql.strip(): lines.append(f"\nSQL:\n {sql.strip()}") # Result table result = ( payload.get("result") or payload.get("resultData") or payload.get("data") ) if isinstance(result, dict): raw_cols = result.get("columns", []) raw_rows = ( result.get("data") or result.get("rows") or result.get("result") or [] ) if raw_cols and raw_rows: # Normalize column spec: may be dicts with title/field or plain strings if raw_cols and isinstance(raw_cols[0], dict): headers = [c.get("title", c.get("field", "?")) for c in raw_cols] col_keys = [c.get("field", "") for c in raw_cols] else: headers = [str(c) for c in raw_cols] col_keys = headers col_widths = [max(len(h), 6) for h in headers] # Calculate widths from data for row in raw_rows[:20]: if isinstance(row, dict): vals = [str(row.get(k, "")) for k in col_keys] elif isinstance(row, (list, tuple)): vals = [str(v) for v in row] else: vals = [str(row)] for i, v in enumerate(vals[:len(col_widths)]): col_widths[i] = max(col_widths[i], len(v)) def _row_str(vals: list) -> str: padded = [str(v).ljust(col_widths[i]) for i, v in enumerate(vals[:len(col_widths)])] return " " + " | ".join(padded) lines.append("") lines.append(_row_str(headers)) lines.append(" " + "-" * (sum(col_widths) + 3 * (len(col_widths) - 1))) for row in raw_rows[:20]: # cap at 20 rows if isinstance(row, dict): vals = [str(row.get(k, "")) for k in col_keys] elif isinstance(row, (list, tuple)): vals = [str(v) for v in row] else: vals = [str(row)] lines.append(_row_str(vals)) if len(raw_rows) > 20: lines.append(f" ... ({len(raw_rows) - 20} more rows)") # If we only produced the header/sep with no real content, suppress content_lines = [l for l in lines if l.strip() and l.strip() != "### Query Result"] return "\n".join(lines) if content_lines else None def _fmt_insights(items: list) -> Optional[str]: """Format a list of insight objects (final conclusion).""" lines = ["\n## Analysis Result\n"] for item in items: title = item.get("title", "") summary = item.get("summary", "") if title: lines.append(f"\n{title}") if summary: lines.append(f" {summary}") # Inline data table if present data_payload = item.get("data") if isinstance(data_payload, str): try: data_payload = json.loads(data_payload) except (json.JSONDecodeError, TypeError): data_payload = None if isinstance(data_payload, dict): columns = data_payload.get("columns", []) rows = data_payload.get("data", []) if columns and rows: lines.append(" " + " | ".join(str(c) for c in columns)) lines.append(" " + "-" * 40) for row in rows[:10]: # cap at 10 rows lines.append(" " + " | ".join(str(v) for v in row)) return "\n".join(lines) def _extract_json_objects(content: str) -> list: """Extract all JSON objects/arrays from a mixed content string. Handles cases like: - Plain text before JSON: "some text {...}" - Multiple JSONs concatenated: '{"a":1}{"b":2}' - Text between JSONs: '{...} text {...}' Returns list of (start, end, parsed_json) tuples. """ results = [] i = 0 while i < len(content): # Find start of JSON object or array if content[i] in ('{', '['): start = i try: parsed, end = json.JSONDecoder().raw_decode(content, i) results.append((start, end, parsed)) i = end except json.JSONDecodeError: i += 1 else: i += 1 return results def _format_data_event(content: str) -> Optional[str]: """Parse and format a data event's content for CLI display. Returns a formatted string to print, or None to suppress the event. """ content = content.strip() if not content: return None # Try to extract all JSON objects/arrays from content json_parts = _extract_json_objects(content) if not json_parts: # No JSON found, return as plain text return content # If entire content is a single JSON, format it if len(json_parts) == 1 and json_parts[0][0] == 0 and json_parts[0][1] == len(content): parsed = json_parts[0][2] return _format_parsed_json(parsed) # Multiple JSONs or mixed text+JSON # Format each JSON part and collect results formatted_parts = [] prev_end = 0 for start, end, parsed in json_parts: # Text before this JSON (if any) if start > prev_end: text_before = content[prev_end:start].strip() if text_before: formatted_parts.append(text_before) # Format the JSON formatted = _format_parsed_json(parsed) if formatted: formatted_parts.append(formatted) prev_end = end # Text after last JSON (if any) if prev_end < len(content): text_after = content[prev_end:].strip() if text_after: formatted_parts.append(text_after) return "\n".join(formatted_parts) if formatted_parts else None def _format_parsed_json(parsed) -> Optional[str]: """Format a parsed JSON value (dict or list) for display.""" # List -> final insights / conclusion if isinstance(parsed, list): if parsed and isinstance(parsed[0], dict) and "title" in parsed[0]: return _fmt_insights(parsed) # Unknown list type - check if it's table summaries if parsed and isinstance(parsed[0], dict) and "table_name" in parsed[0]: return _fmt_table_summaries(parsed) return None # suppress other lists # Dict if isinstance(parsed, dict): result_type = parsed.get("result_type") if result_type == "jupyter_cell": return _fmt_jupyter_cell(parsed) return None # other tool metadata, suppress return None def _fmt_table_summaries(items: list) -> Optional[str]: """Format table summary list for display.""" lines = ["\n### Table Summaries\n"] for item in items: table_name = item.get("table_name", "?") summary = item.get("table_summary", "") lines.append(f"\n{table_name}:") if summary: # Wrap long summaries words = summary.split() line = " " for word in words: if len(line) + len(word) + 1 > 70: lines.append(line) line = " " + word else: line += (" " if line != " " else "") + word if line.strip(): lines.append(line) lines.append(sep) return "\n".join(lines) def _fmt_plan_progress(plan_data: dict, current_step: int, total_steps: int) -> Optional[str]: """Format a plan data event showing step progress. Expected plan_data structure (from API): {"current_step": 2, "plan_status": "...", "plans": [{"plan": {"steps": [...]}}]} Each step has "name" and optionally "description". """ steps = [] plans = plan_data.get("plans", []) if plans and isinstance(plans[0], dict): inner_plan = plans[0].get("plan", {}) if isinstance(inner_plan, dict): steps = inner_plan.get("steps", []) step_name = "" step_desc = "" # current_step is 1-based from API if steps and 0 < current_step <= len(steps): step = steps[current_step - 1] step_name = step.get("name", "") step_desc = step.get("description", "") label = f"[Plan] Step {current_step}/{total_steps}" if step_name: label += f": {step_name}" lines = [label] if step_desc: lines.append(f" {step_desc}") return "\n".join(lines) def _fmt_status_change(content: str) -> Optional[str]: """Format a status_change event's JSON content. Expected JSON: {"previous": "PLANNING", "current": "STEP_EXECUTION", "current_task": "..."} """ try: data = json.loads(content) if isinstance(content, str) else content except (json.JSONDecodeError, TypeError): return None previous = data.get("previous", "none") current = data.get("current", "none") current_task = data.get("current_task", "") line = f"[Task] {previous} \u2192 {current}" if current_task and current_task != current: line += f" ({current_task})" return line def _fmt_output_conclusion(text: str, output_dir: Optional[Path] = None, header: Optional[str] = None) -> str: """Wrap analysis conclusion text with visual borders. If output_dir is provided, extracts inline base64 images from the text, saves them to output_dir/images/, and replaces inline data with file paths. If header is provided, use it instead of the default "Analysis Conclusion". """ if output_dir is not None: text = _extract_and_save_images(text, output_dir) title = header or "Analysis Conclusion" return f"\n## {title}\n\n{text}\n" # Regex for markdown inline images with base64 data URIs: # ![alt](data:image/png;base64,AAAA...) _BASE64_IMG_RE = re.compile( r"!\[([^\]]*)\]" # ![alt text] r"\(" # ( r"data:image/([a-zA-Z0-9.+-]+)" # data:image/png r";base64," # ;base64, r"([A-Za-z0-9+/=\s]+)" # base64 payload (may contain whitespace) r"\s*\)" # ) ) def _extract_and_save_images(text: str, output_dir: Path, prefix: str = "conclusion") -> str: """Extract inline base64 images from text and save to disk. Returns cleaned text with images replaced by file-path references. """ img_dir = output_dir / "images" counter = 0 def _replace(match: re.Match) -> str: nonlocal counter counter += 1 alt = match.group(1) or f"image_{counter}" mime_subtype = match.group(2).lower() # png, jpeg, svg+xml, etc. b64_data = match.group(3).strip() # Determine file extension ext = _mime_subtype_to_ext(mime_subtype) filename = f"{prefix}_{counter}{ext}" save_path = img_dir / filename try: img_dir.mkdir(parents=True, exist_ok=True) raw = base64.b64decode(b64_data) save_path.write_bytes(raw) return f"![{alt}]({save_path.resolve()})" except Exception: return f"![{alt}](save_failed)" return _BASE64_IMG_RE.sub(_replace, text) def _mime_subtype_to_ext(subtype: str) -> str: """Map image MIME subtype to file extension.""" mapping = { "png": ".png", "jpeg": ".jpg", "jpg": ".jpg", "gif": ".gif", "svg+xml": ".svg", "webp": ".webp", "bmp": ".bmp", "tiff": ".tiff", } return mapping.get(subtype, f".{subtype.split('+')[0]}") def _fmt_recommended_questions(content: str) -> Optional[str]: """Format recommended follow-up questions. Expected JSON: {"questions": ["q1", "q2", ...]} or a list of strings. """ try: data = json.loads(content) if isinstance(content, str) else content except (json.JSONDecodeError, TypeError): return None questions = [] if isinstance(data, dict): questions = data.get("questions", []) if not questions: questions = data.get("recommendQuestions", []) elif isinstance(data, list): questions = data if not questions: return None lines = ["\n### Suggestions"] for i, q in enumerate(questions[:5], 1): text = q.get("question", q) if isinstance(q, dict) else str(q) lines.append(f" {i}. {text}") return "\n".join(lines) def _fmt_ask_report_render(content: str, session_id: Optional[str] = None) -> Optional[str]: """Format ask_report_render confirmation prompt.""" lines = ["\n### Report Render"] # Try to parse as JSON first in case it has a 'message' field display_text = content try: data = json.loads(content) if isinstance(content, str) else content if isinstance(data, dict) and "message" in data: display_text = data["message"] except (json.JSONDecodeError, TypeError): pass preview = display_text[:500] + ("..." if len(display_text) > 500 else "") lines.append(f" {preview}") sid = session_id or "" lines.append(f"\n> ⚠️ Please review the report rendering request.") lines.append(f"> To confirm you want to render the report, DO NOT create a new session, use the existing session and explicitly say yes:") lines.append(f"> python3 scripts/data_agent_cli.py attach --session-id {sid} -q '确认绘制网页报告'") lines.append(f"> python3 scripts/data_agent_cli.py attach --session-id {sid} -q 'yes, render the report'") lines.append(f"> To decline, simply respond with no:") lines.append(f"> python3 scripts/data_agent_cli.py attach --session-id {sid} -q '不绘制'") return "\n".join(lines) FILE:scripts/cli/log_handler.py """Structured logging handler for Data Agent. Provides logging to plain text format only (progress.jsonl disabled). """ import json from pathlib import Path from datetime import datetime from typing import TextIO, Optional, Dict, Any class StructuredLogHandler: """Handles logging to plain text format only (JSONL logging disabled).""" def __init__(self, session_dir: Path, log_prefix: str = "progress"): """Initialize the log handler (JSONL logging disabled). Args: session_dir: Directory for the session where logs will be stored. log_prefix: Prefix for log files (e.g., "progress" for progress.log, or "process" for process.log) """ self.session_dir = session_dir self.log_prefix = log_prefix self.progress_log_path = session_dir / f"{log_prefix}.log" # JSONL logging disabled - progress.jsonl creation is skipped # self.progress_jsonl_path = session_dir / f"{log_prefix}.jsonl" # Open only the plain text log file self.progress_log_file: Optional[TextIO] = None # JSONL logging disabled # self.progress_jsonl_file: Optional[TextIO] = None def __enter__(self): """Enter context manager, opening log file.""" # Ensure directory exists self.session_dir.mkdir(parents=True, exist_ok=True) # Open only the plain text log file self.progress_log_file = open(self.progress_log_path, "w", encoding="utf-8") # JSONL logging disabled return self def __exit__(self, exc_type, exc_val, exc_tb): """Exit context manager, closing log file.""" if self.progress_log_file: self.progress_log_file.close() # JSONL logging disabled def write_log(self, text: str): """Write text to the plain log file (e.g., progress.log or process.log). Args: text: Text to write to the plain text log. """ if self.progress_log_file: self.progress_log_file.write(text) self.progress_log_file.flush() # Ensure immediate write def write_jsonl(self, data: Dict[str, Any]): """Write structured data to the JSONL file (disabled). Args: data: Dictionary containing structured log data. """ # JSONL logging is disabled - do nothing pass def write_both(self, text: str, data: Optional[Dict[str, Any]] = None): """Write to plain text log only (JSONL logging disabled). Args: text: Text to write to the plain text log. data: Structured data (ignored since JSONL logging is disabled). """ self.write_log(text) # JSONL logging disabled - data parameter is ignored def close(self): """Close the plain text log file handle (JSONL logging disabled).""" if self.progress_log_file: self.progress_log_file.close() self.progress_log_file = None # JSONL logging disabled FILE:scripts/cli/notify.py """Notification utilities for pushing updates to Clawdbot / OpenClaw sessions. Author: Tinker Created: 2026-03-11 """ import os import json import subprocess from urllib.request import Request, urlopen from urllib.error import URLError def get_active_session() -> str: """Get the active OpenClaw / Clawdbot session ID.""" print("[Notify] Trying to determine active session...", flush=True) # 1. Environment variables session = os.environ.get("OPENCLAW_SESSION") or os.environ.get("CLAWDBOT_PUSH_SESSION") if session: print(f"[Notify] Found session from environment: {session}", flush=True) return session # 2. Try CLI cli_cmd = None for cmd in ["openclaw", "clawdbot"]: try: subprocess.run(["which", cmd], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) cli_cmd = cmd print(f"[Notify] Found CLI: {cli_cmd}", flush=True) break except subprocess.CalledProcessError: continue if not cli_cmd: print("[Notify] No CLI found (openclaw or clawdbot).", flush=True) return "" try: print(f"[Notify] Running: {cli_cmd} sessions --active 5 --json", flush=True) result = subprocess.run( [cli_cmd, "sessions", "--active", "5", "--json"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=5 ) if result.returncode == 0: data = json.loads(result.stdout) sessions = data.get("sessions", []) if sessions and sessions[0].get("key"): key = sessions[0]["key"] # Handle "webchat:xxx" format session_key = key.split(":")[-1] if ":" in key else key print(f"[Notify] Found session from CLI: {session_key}", flush=True) return session_key else: print("[Notify] CLI returned success but no active sessions found.", flush=True) else: print(f"[Notify] CLI command failed. Return code: {result.returncode}, Stderr: {result.stderr}", flush=True) except Exception as e: print(f"[Notify] Exception running CLI to get session: {e}", flush=True) return "" def push_notification(session_id: str, message: str) -> bool: """Push a notification message to the specified session or active session. Args: session_id: The target session ID (optional). message: The message content to push. Returns: True if successful, False otherwise. """ print(f"\n[Notify] Starting push notification... Target session_id: {session_id}", flush=True) custom_url = os.environ.get("ASYNC_TASK_PUSH_URL") # 1. Custom HTTP push if custom_url: print(f"[Notify] ASYNC_TASK_PUSH_URL is set: {custom_url}", flush=True) token = os.environ.get("ASYNC_TASK_AUTH_TOKEN", "") data = json.dumps({ "sessionId": session_id, "content": message, "role": "assistant" }).encode("utf-8") headers = { "Content-Type": "application/json", "Content-Length": str(len(data)) } if token: headers["Authorization"] = f"Bearer {token}" req = Request(custom_url, data=data, headers=headers, method="POST") try: with urlopen(req, timeout=10) as response: success = response.status >= 200 and response.status < 300 print(f"[Notify] HTTP push result: status={response.status}, success={success}", flush=True) return success except URLError as e: print(f"[Notify] HTTP push failed with URLError: {e}", flush=True) return False # 2. CLI push print("[Notify] Falling back to CLI push...", flush=True) target_session = session_id or get_active_session() if not target_session: print("[Notify] Aborting push: No target session available.", flush=True) return False cli_cmd = None for cmd in ["openclaw", "clawdbot"]: try: subprocess.run(["which", cmd], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) cli_cmd = cmd print(f"[Notify] Found CLI for push: {cli_cmd}", flush=True) break except subprocess.CalledProcessError: continue if not cli_cmd: print("[Notify] Aborting push: No CLI found.", flush=True) return False cmd_args = [cli_cmd, "sessions", "send", "--session", target_session, message] print(f"[Notify] Executing: {' '.join(cmd_args)}", flush=True) try: result = subprocess.run( cmd_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=30, text=True ) success = result.returncode == 0 if success: print("[Notify] CLI push succeeded.", flush=True) else: print(f"[Notify] CLI push failed. Return code: {result.returncode}, Stderr: {result.stderr}", flush=True) return success except Exception as e: print(f"[Notify] Exception running CLI push: {e}", flush=True) return False FILE:scripts/cli/parser.py """Argument parsing and CLI entry point. Author: Tinker Created: 2026-03-04 """ import argparse import os import sys from cli.cmd_db import cmd_db from cli.cmd_ls import cmd_ls from cli.cmd_file import cmd_file from cli.cmd_attach import cmd_attach from cli.cmd_dms import cmd_dms from cli.cmd_import import cmd_import from cli.cmd_reports import cmd_reports def build_parser() -> argparse.ArgumentParser: """Build command-line argument parser.""" parser = argparse.ArgumentParser( prog="data_agent_cli", description="Data Agent Unified CLI - Database Query/Analysis & File Analysis", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Database Mode Description ASK_DATA: Quick query mode with SQL execution + natural language response ANALYSIS: Deep analysis mode with planning, multi-step reasoning, and report generation Examples: # -- Database ASK_DATA (default) -- python data_agent_cli.py db \\ --dms-instance-id --dms-db-id \\ --instance-name --db-name chinook \\ --tables "album,artist,customer" \\ -q "Who has the highest sales?" # -- Database ANALYSIS mode -- python data_agent_cli.py db --session-mode ANALYSIS \\ --dms-instance-id --dms-db-id \\ --instance-name --db-name chinook \\ --tables "album,artist,customer" \\ -q "Analyze sales trends and generate report" # -- File Analysis -- python data_agent_cli.py file /path/to/data.csv python data_agent_cli.py file /path/to/data.csv -q "Analyze sales trends" """, ) subparsers = parser.add_subparsers(dest="command", metavar="COMMAND") subparsers.required = True # -- db subcommand -- db_parser = subparsers.add_parser( "db", help="Database query/analysis (supports ASK_DATA and ANALYSIS modes)", description="""Connect to DMS database for natural language queries or data analysis. Mode Description: ASK_DATA (default): Quick SQL query + natural language response ANALYSIS : Deep analysis, multi-step reasoning, report generation""", formatter_class=argparse.RawDescriptionHelpFormatter, ) # Data source parameters ds_group = db_parser.add_argument_group("Data Source Parameters (required)") ds_group.add_argument( "--dms-instance-id", type=int, metavar="INT", help="DMS Instance ID (e.g., 1234567)", ) ds_group.add_argument( "--dms-db-id", type=int, metavar="INT", help="DMS Database ID (e.g., 12345678)", ) ds_group.add_argument( "--instance-name", metavar="STR", help="RDS instance name (e.g., rm-xxxxxx.mysql.rds.aliyuncs.com)", ) ds_group.add_argument( "--db-name", metavar="STR", help="Database name (e.g., chinook)", ) ds_group.add_argument( "--tables", metavar="T1,T2,...", help="Table names to include (comma-separated, e.g., album,artist,track)", ) ds_group.add_argument( "--table-ids", metavar="ID1,ID2,...", help="Data Center table IDs (comma-separated)", ) ds_group.add_argument( "--engine", default="mysql", metavar="ENGINE", help="Database engine type (default: mysql)", ) # Default region from env, fallback to cn-hangzhou _default_region = os.environ.get("DATA_AGENT_REGION", "cn-hangzhou") ds_group.add_argument( "--region", default=_default_region, metavar="REGION", help=f"Region ID (default: {_default_region})", ) # Query db_parser.add_argument( "--query", "-q", metavar="QUERY", help="Execute a single natural language query (runs preset queries if omitted)", ) db_parser.add_argument( "--async-run", action=argparse.BooleanOptionalAction, default=True, help="Run asynchronously in background (default: True). Use --no-async-run for synchronous mode.", ) # Session options db_parser.add_argument( "--session-mode", default="ASK_DATA", choices=["ASK_DATA", "ANALYSIS", "INSIGHT"], metavar="MODE", help="Session mode: ASK_DATA (default) | ANALYSIS | INSIGHT", ) db_parser.add_argument( "--output", choices=["summary", "detail", "raw"], default="summary", metavar="MODE", help="Output verbosity: summary (default, conclusion only) | detail (full process) | raw (every SSE event)", ) db_parser.add_argument( "--enable-search", action="store_true", help="Enable search capability in the session (default: False)", ) db_parser.set_defaults(func=cmd_db) # -- file subcommand -- file_parser = subparsers.add_parser( "file", help="File upload and analysis (CSV / Excel / JSON)", description="Upload local data files for intelligent analysis via Data Agent.", ) # Create a mutually exclusive group for file source file_source_group = file_parser.add_mutually_exclusive_group(required=True) file_source_group.add_argument( "file_path", nargs='?', # Make it optional since we have --file-id metavar="FILE", help="Local file path to upload for analysis (supports: .csv / .xlsx / .xls / .json / .txt)", ) file_source_group.add_argument( "--file-id", metavar="FILE_ID", help="Data Center file ID to analyze directly (e.g., f-8941bx83xy9513xvpewrha01m)", ) # Query file_parser.add_argument( "--query", "-q", metavar="QUERY", help="Execute a single custom query (uses preset questions if not specified)", ) file_parser.add_argument( "--async-run", action=argparse.BooleanOptionalAction, default=True, help="Run asynchronously in background (default: True). Use --no-async-run for synchronous mode.", ) file_parser.add_argument( "--list-generated-files", action="store_true", help="List agent-generated reports/charts after analysis", ) file_parser.add_argument( "--session-mode", default="ANALYSIS", choices=["ASK_DATA", "ANALYSIS", "INSIGHT"], metavar="MODE", help="Session mode: ASK_DATA | ANALYSIS (default) | INSIGHT", ) file_parser.add_argument( "--output", default="summary", choices=["summary", "detail", "raw"], metavar="MODE", help="Output mode: summary (default, conclusion only) | detail (full process) | raw (all SSE events)", ) file_parser.add_argument( "--enable-search", action="store_true", help="Enable search capability in the session (default: False)", ) file_parser.set_defaults(func=cmd_file) # -- attach subcommand -- attach_parser = subparsers.add_parser( "attach", help="Connect to an existing session to continue analysis or confirm plan", description=( "Attach to an existing Data Agent session. Use this to:\n" " - Continue an interrupted ANALYSIS session\n" " - Confirm an execution plan after WAIT_INPUT\n" " - Ask follow-up questions in an existing session\n\n" "Examples:\n" " python3 scripts/data_agent_cli.py attach --session-id -q '\u786e\u8ba4\u6267\u884c\u8ba1\u5212'\n" " python3 scripts/data_agent_cli.py attach --session-id " ), formatter_class=argparse.RawDescriptionHelpFormatter, ) attach_parser.add_argument( "--session-id", required=True, metavar="SESSION_ID", help="Session ID to connect to", ) attach_parser.add_argument( "--output", choices=["summary", "detail", "raw"], default="summary", metavar="MODE", help="Output verbosity: summary (default, conclusion only) | detail (full process) | raw (every SSE event)", ) attach_parser.add_argument( "--checkpoint", type=int, metavar="CHECKPOINT", help="Specific checkpoint to resume from (e.g. 52). Resumes from breakpoint.", ) attach_parser.add_argument( "--from-start", action="store_true", help="Replay from checkpoint=0 (full history) instead of watching live stream", ) attach_parser.add_argument( "--query", "-q", metavar="QUERY", help="Execute a single query (e.g., '\u786e\u8ba4\u6267\u884c\u8ba1\u5212' or '\u7ee7\u7eed\u5206\u6790')", ) attach_parser.add_argument( "--async-run", action=argparse.BooleanOptionalAction, default=True, help="Run asynchronously in background (default: True). Use --no-async-run for synchronous mode.", ) attach_parser.set_defaults(func=cmd_attach) # -- ls subcommand -- ls_parser = subparsers.add_parser( "ls", help="List DMS databases and tables with analysis-ready parameters", description=( "List databases registered in DMS Data Center and print the\n" "key IDs and parameters needed for the db subcommand.\n\n" "Examples:\n" " python3 scripts/data_agent_cli.py ls # all databases\n" " python3 scripts/data_agent_cli.py ls --search chinook # filter by name\n" " python3 scripts/data_agent_cli.py ls --db-id # tables in DB" ), formatter_class=argparse.RawDescriptionHelpFormatter, ) ls_parser.add_argument( "--search", "-s", metavar="KEYWORD", help="Filter databases by keyword (database or instance name)", ) ls_parser.add_argument( "--db-id", metavar="DB_ID", help="DbId of a specific database; lists its tables and prints the ready-to-use db command", ) ls_parser.set_defaults(func=cmd_ls) # -- dms subcommand -- dms_parser = subparsers.add_parser( "dms", help="DMS tools integration (list-instances, search-database, list-tables)", description="""Access DMS Server tools for instance, database and table management. Tools: list-instances Search and list database instances in DMS search-database Search databases by schema name list-tables List tables in a database Examples: python3 scripts/data_agent_cli.py dms list-instances python3 scripts/data_agent_cli.py dms list-instances --search mysql --db-type mysql python3 scripts/data_agent_cli.py dms search-database --search-key mydb python3 scripts/data_agent_cli.py dms list-tables --database-id """, formatter_class=argparse.RawDescriptionHelpFormatter, ) dms_parser.add_argument( "tool", choices=["list-instances", "search-database", "list-tables"], help="DMS tool to execute", ) # list-instances options dms_parser.add_argument( "--search", metavar="KEYWORD", help="Search key for instances (host, alias, etc.)", ) dms_parser.add_argument( "--db-type", metavar="TYPE", help="Database type filter (e.g., mysql, polardb, oracle)", ) dms_parser.add_argument( "--env-type", metavar="ENV", help="Environment type filter (e.g., product, dev, test)", ) # search-database options dms_parser.add_argument( "--search-key", metavar="KEYWORD", help="Schema name search keyword for search-database", ) dms_parser.add_argument( "--page-number", type=int, default=1, metavar="PAGE", help="Page number for pagination (default: 1)", ) dms_parser.add_argument( "--page-size", type=int, default=200, metavar="SIZE", help="Page size for pagination (default: 200, max: 1000)", ) # list-tables options dms_parser.add_argument( "--database-id", metavar="DB_ID", help="Database ID for list-tables", ) dms_parser.add_argument( "--search-name", metavar="NAME", help="Table name search keyword for list-tables", ) dms_parser.set_defaults(func=cmd_dms) # -- import subcommand -- import_parser = subparsers.add_parser( "import", help="Import DMS database tables to Data Agent Data Center", description="""Add DMS database tables to Data Agent Data Center. This command imports tables from DMS into Data Agent's Data Center, making them available for analysis via the 'db' subcommand. Examples: python3 scripts/data_agent_cli.py import \\ --dms-instance-id \\ --dms-db-id \\ --instance-name \\ --db-name employees \\ --tables "departments,employees,salaries" """, formatter_class=argparse.RawDescriptionHelpFormatter, ) # Data source parameters (required) import_group = import_parser.add_argument_group("Data Source Parameters (required)") import_group.add_argument( "--dms-instance-id", type=int, metavar="INT", help="DMS Instance ID (e.g., 1234567)", ) import_group.add_argument( "--dms-db-id", type=int, metavar="INT", help="DMS Database ID (e.g., 12345678)", ) import_group.add_argument( "--instance-name", metavar="STR", help="RDS instance name (e.g., rm-xxxxxx.mysql.rds.aliyuncs.com)", ) import_group.add_argument( "--db-name", metavar="STR", help="Database name (e.g., employees)", ) import_group.add_argument( "--tables", metavar="T1,T2,...", help="Table names to import (comma-separated, e.g., departments,employees,salaries)", ) import_group.add_argument( "--engine", default="mysql", metavar="ENGINE", help="Database engine type (default: mysql)", ) _default_region = os.environ.get("DATA_AGENT_REGION", "cn-hangzhou") import_group.add_argument( "--region", default=_default_region, metavar="REGION", help=f"Region ID (default: {_default_region})", ) import_parser.set_defaults(func=cmd_import) # -- reports subcommand -- reports_parser = subparsers.add_parser( "reports", help="List and download agent-generated files (reports, charts, data files)", description="""List and download files generated during an ANALYSIS or INSIGHT session. Examples: python3 scripts/data_agent_cli.py reports --session-id """, formatter_class=argparse.RawDescriptionHelpFormatter, ) reports_parser.add_argument( "--session-id", required=True, metavar="SESSION_ID", help="Session ID to fetch reports for", ) # The download parameter is no longer an optional flag; # reports command now downloads files automatically. # We keep the argument parser clean by removing it. reports_parser.set_defaults(func=cmd_reports) return parser def main() -> None: parser = build_parser() args = parser.parse_args() try: args.func(args) except KeyboardInterrupt: print("\n\nInterrupted.", file=sys.stderr) sys.exit(130) except Exception as e: print(f"\nError: {e}", file=sys.stderr) sys.exit(1) FILE:scripts/cli/streaming.py """SSE event routing and streaming response handling. Uses a stateful StreamState + event_type->category dispatch architecture, modelled after the TypeScript transform-chat-message-chunk.ts reference. Output modes: "summary" (default) -- minimal: only conclusion, query result, confirmations, suggestions, reports. Process details suppressed. "detail" -- full: plan progress, status changes, step titles, jupyter cell results, plus everything in summary. "raw" -- every SSE event printed verbatim. Author: Tinker Created: 2026-03-03 """ import json from dataclasses import dataclass, field, asdict from pathlib import Path from typing import Optional from datetime import datetime import sys import threading from cli.log_handler import StructuredLogHandler # --------------------------------------------------------------------------- # StreamState — shared mutable state across events within one stream # --------------------------------------------------------------------------- @dataclass class StreamState: """Mutable state shared across events within a single SSE stream.""" # Output mode: "summary" (minimal) | "detail" (full) | "raw" output_mode: str = "summary" # Accumulated textual output for capturing the final result full_output: list[str] = field(default_factory=list) # Output directory for saving extracted images etc. output_dir: Optional[Path] = None # Log handler for process logs (if provided) process_log_handler: Optional[StructuredLogHandler] = None # Session metadata for structured logging session_id: Optional[str] = None session_status: Optional[str] = None # Last received checkpoint (for checkpoint.txt generation) last_checkpoint: Optional[int] = None # Cross-event accumulators for content groups content_category: Optional[str] = None pending_step_label: Optional[str] = None output_conclusion_chunks: Optional[list[str]] = None tool_call_response_content: Optional[str] = None ask_report_render_payload: Optional[dict] = None # Stream completion state got_content: bool = False need_user_confirm: bool = False is_attach: bool = False # True when called from attach command from cli.formatters import ( _SKIP_DATA_CATEGORIES, _extract_json_objects, _fmt_insights, _fmt_jupyter_cell, _fmt_output_conclusion, _fmt_plan_progress, _fmt_recommended_questions, _fmt_status_change, _fmt_task_finish, _fmt_ask_report_render, ) from data_agent import SSEEvent, MessageHandler, DataSource # Global variables for structured logging _progress_jsonl_file = None _progress_log_file = None # New variable for plain text log _jsonl_lock = threading.Lock() def init_structured_logging(output_dir: Optional[Path] = None): """Initialize structured logging to plain text format only (progress.jsonl disabled). Args: output_dir: Directory where progress.log should be created """ global _progress_jsonl_file, _progress_log_file if output_dir: # Disabled JSONL logging - progress.jsonl creation is skipped _progress_jsonl_file = None # Only open progress.log if not already set if not _progress_log_file: log_path = output_dir / "progress.log" # In worker processes, stdout is redirected to progress.log, # so check if that's the case and reuse that fd if hasattr(sys.stdout, 'name') and sys.stdout.name == str(log_path): # stdout is already redirected to progress.log, reuse that fd _progress_log_file = sys.stdout else: # In sync mode, open the log file directly _progress_log_file = open(log_path, "w", encoding="utf-8") def close_structured_logging(): """Close the structured logging files.""" global _progress_jsonl_file, _progress_log_file if _progress_jsonl_file: _progress_jsonl_file.close() _progress_jsonl_file = None # Only close _progress_log_file if it's not sys.stdout (which happens # when stdout is redirected to progress.log in worker processes) if _progress_log_file and _progress_log_file is not sys.stdout: _progress_log_file.close() _progress_log_file = None def write_to_jsonl(data: dict): """Write structured data to JSONL file. Args: data: Dictionary containing structured log data """ # JSONL logging is disabled - do nothing pass def write_to_progress_log(text: str): """Write text to progress log file. Args: text: Text to write to the plain text log """ global _progress_log_file # Skip writing if in worker process where stdout is redirected to progress.log # This prevents duplication since print() calls already go to the same file import os if os.environ.get("DATA_AGENT_ASYNC_WORKER") == "1": return if _progress_log_file: _progress_log_file.write(text) _progress_log_file.flush() else: # Debug: log to file if _progress_log_file is None debug_file = Path("/tmp/data_agent_debug.log") with open(debug_file, "a") as f: f.write(f"DEBUG: _progress_log_file is None, cannot write: {text[:50]}...\n") def _out(state: StreamState, text: str = "", **kwargs) -> None: """Print text to console and record it as result output. Use for actual Data Agent results that should appear in output.md: conclusions, plans, SQL, reports, insights, recommendations, etc. """ # In worker process, force flush after each print for real-time logging import os if os.environ.get("DATA_AGENT_ASYNC_WORKER") == "1": kwargs['flush'] = True print(text, **kwargs) state.full_output.append(text) # Write to progress.log (for both sync and async modes) # Skip if in worker process where stdout is redirected to progress.log # to prevent duplication if os.environ.get("DATA_AGENT_ASYNC_WORKER") != "1" and text.strip(): progress_text = text + ("\n" if kwargs.get('end', '\n') == '\n' else "") write_to_progress_log(progress_text) # Write structured log entry to JSONL if text.strip(): # Only log non-empty text write_to_jsonl({ 'type': 'output', 'content': text, 'category': 'result' }) # Also write to process logs if handler is available if state.process_log_handler and text.strip(): state.process_log_handler.write_both(text + ("\n" if kwargs.get('end', '\n') == '\n' else "")) def _log(state: StreamState, text: str = "", **kwargs) -> None: """Print text to console only (progress.log) without recording to output.md. Use for diagnostic / process information: user query echo, plan step progress, status changes, intermediate code execution results, etc. """ # In worker process, force flush after each print for real-time logging import os if os.environ.get("DATA_AGENT_ASYNC_WORKER") == "1": kwargs['flush'] = True print(text, **kwargs) # Write to progress.log (for both sync and async modes) # Skip if in worker process where stdout is redirected to progress.log # to prevent duplication if os.environ.get("DATA_AGENT_ASYNC_WORKER") != "1" and text.strip(): progress_text = text + ("\n" if kwargs.get('end', '\n') == '\n' else "") write_to_progress_log(progress_text) # Write structured log entry to JSONL if text.strip(): # Only log non-empty text write_to_jsonl({ 'type': 'log', 'content': text, 'category': 'diagnostic' }) # Also write to process logs if handler is available if state.process_log_handler and text.strip(): state.process_log_handler.write_both(text + ("\n" if kwargs.get('end', '\n') == '\n' else "")) # --------------------------------------------------------------------------- # StreamState — shared mutable state across events within one stream # --------------------------------------------------------------------------- @dataclass class StreamState: """Mutable state shared across events within a single SSE stream.""" # Output mode: "summary" (minimal) | "detail" (full) | "raw" output_mode: str = "summary" # Accumulated textual output for capturing the final result full_output: list[str] = field(default_factory=list) # Output directory for saving extracted images etc. output_dir: Optional[Path] = None # Log handler for process logs (if provided) process_log_handler: Optional[StructuredLogHandler] = None # Session ID for user prompts session_id: Optional[str] = None session_status: Optional[str] = None is_attach: bool = False # Task / Plan tracking current_task: Optional[str] = None current_step: Optional[int] = None current_plan_status: Optional[str] = None total_steps: int = 0 step_names: dict = field(default_factory=dict) # {order: name} # Pending step label — consumed by next conclusion output (summary mode) pending_step_label: Optional[str] = None # Current user question (extracted from chat_start) current_question: Optional[str] = None # content_start / content_finish lifecycle content_category: Optional[str] = None # Accumulators (between content_start and content_finish) output_conclusion_chunks: Optional[list] = None tool_call_response_content: str = "" # Result flags got_content: bool = False need_user_confirm: bool = False # Turn counter for multi-turn display turn_count: int = 0 # Store current checkpoint for resumption last_checkpoint: Optional[int] = None # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def _print_event( event: SSEEvent, output_mode: str = "summary", state: Optional[StreamState] = None, ) -> tuple: """Route a single SSE event to the appropriate output handler. Args: event: The SSE event to handle. output_mode: "summary" (minimal), "detail" (full), or "raw". state: Optional shared StreamState for cross-event tracking. When None a temporary state is created (backward compat). Returns: (got_content: bool, need_user_confirm: bool) """ content = event.content or "" if state is not None and event.checkpoint is not None: state.last_checkpoint = event.checkpoint if state.output_dir: try: with open(state.output_dir / "checkpoint.txt", "w") as f: f.write(str(event.checkpoint)) except Exception: pass # -- raw mode -- (diagnostic, not recorded to output.md) if output_mode == "raw": cat = f" cat={event.category}" if event.category else "" cp = f" cp={event.checkpoint}" if event.checkpoint is not None else "" ct = f" ct={event.content_type}" if getattr(event, "content_type", None) else "" msg1 = f"[{event.event_type}]{cat}{cp}{ct}" if state is not None: _log(state, msg1) else: print(msg1) if content: preview = content[:600] + ("...(truncated)" if len(content) > 600 else "") msg2 = f" {preview}" if state is not None: _log(state, msg2) else: print(msg2) return bool(content), False # -- summary / detail: delegate to stateful dispatch -- if state is None: state = StreamState(output_mode=output_mode) _dispatch_event(state, event) return state.got_content, state.need_user_confirm def _stream_response( message_handler: MessageHandler, session, query: str, data_source: Optional[DataSource] = None, output_mode: str = "summary", output_dir: Optional[Path] = None, process_log_handler: Optional[StructuredLogHandler] = None, is_attach: bool = False, ) -> tuple[bool, bool, str]: """Stream response tokens to stdout in real-time. Uses StreamState to track cross-event context (plan progress, output_conclusion accumulation, tool_call_response lifecycle). Returns (got_content: bool, need_user_confirm: bool, output_text: str). """ state = StreamState(output_mode=output_mode) state.output_dir = output_dir state.process_log_handler = process_log_handler # Set the process log handler state.session_id = getattr(session, 'session_id', None) state.session_status = getattr(session, 'status', None) if hasattr(state.session_status, 'value'): state.session_status = state.session_status.value # IMPORTANT: Set is_attach to True for attach mode state.is_attach = is_attach # Initialize structured logging if output directory provided init_structured_logging(output_dir) try: for event in message_handler.stream_events(session, query, data_source=data_source): # Log SSE event to JSONL (disabled per user request to disable progress.jsonl) # write_to_jsonl({ # 'type': 'sse_event', # 'event_type': event.event_type, # 'category': event.category, # 'content': event.content, # 'data': event.data, # 'raw_event': asdict(event) # }) if event.event_type == "SSE_FINISH": break _print_event(event, output_mode, state=state) finally: # Close structured logging close_structured_logging() _finalize_stream(state) if state.got_content and not state.need_user_confirm: _out(state, "") # final newline after streaming full_text = "\n".join(state.full_output) return state.got_content, state.need_user_confirm, full_text def _finalize_stream(state: StreamState) -> None: """Flush any pending accumulators when the stream ends.""" # Flush accumulated output_conclusion that never got a content_finish if state.output_conclusion_chunks: text = "".join(state.output_conclusion_chunks) if text.strip(): header = _consume_step_label(state) _out(state, _fmt_output_conclusion(text, output_dir=state.output_dir, header=header)) state.got_content = True state.output_conclusion_chunks = None # Flush accumulated tool_call_response that never got a content_finish if state.tool_call_response_content: _flush_tool_call_response(state) state.content_category = None def _consume_step_label(state: StreamState) -> Optional[str]: """Return and clear the pending step label for use as conclusion header.""" label = state.pending_step_label state.pending_step_label = None return label def _is_user_confirmation_event(event: SSEEvent) -> bool: """Check if an SSE event requires user confirmation.""" if event.event_type == "data": data = event.data or {} if data.get("need_confirm") or data.get("requires_confirmation"): return True category = event.category or "" if "confirm" in category.lower() or "approval" in category.lower(): return True content = event.content or "" confirm_keywords = ["确认", "confirm", "approve", "approval", "需要确认"] if any(keyword in content.lower() for keyword in confirm_keywords): return True return False # --------------------------------------------------------------------------- # Dispatch — event_type first-level switch # --------------------------------------------------------------------------- def _dispatch_event(state: StreamState, event: SSEEvent) -> None: """Main dispatcher: routes by event_type to specialised handlers.""" et = event.event_type if et == "chat_start": _handle_chat_start(state, event) elif et == "content_start": _handle_content_start(state, event) elif et == "delta": _handle_delta(state, event) elif et == "data": _handle_data(state, event) elif et == "content_finish": _handle_content_finish(state, event) elif et == "status_change": _handle_status_change(state, event) elif et == "chat_finish": _handle_chat_finish(state, event) elif et == "chat_canceled": _out(state, "\n[Canceled] 任务已取消") state.got_content = True elif et == "SSE_FAILURE": _handle_sse_failure(state, event) # --------------------------------------------------------------------------- # Handlers # --------------------------------------------------------------------------- def _handle_chat_start(state: StreamState, event: SSEEvent) -> None: """Handle chat_start events — mark the beginning of a new turn.""" state.turn_count += 1 # Reset per-turn plan state state.current_step = None state.current_plan_status = None state.pending_step_label = None content = event.content or "" # Extract user question from JSON payload question = "" try: data = json.loads(content) if isinstance(content, str) else content if isinstance(data, dict): question = data.get("message", "") except (json.JSONDecodeError, TypeError): pass if question: state.current_question = question if question: _log(state, f"\n> User Query: {question}\n") def _handle_content_start(state: StreamState, event: SSEEvent) -> None: """Set content_category and initialise accumulators.""" # If previous content never finished, flush it if state.content_category == "tool_call_response" and state.tool_call_response_content: _flush_tool_call_response(state) if state.content_category == "output_conclusion" and state.output_conclusion_chunks: text = "".join(state.output_conclusion_chunks) if text.strip(): header = _consume_step_label(state) _out(state, _fmt_output_conclusion(text, output_dir=state.output_dir, header=header)) state.got_content = True state.output_conclusion_chunks = None state.content_category = event.category if event.category == "tool_call_response": state.tool_call_response_content = "" elif event.category == "output_conclusion": state.output_conclusion_chunks = [] def _handle_delta(state: StreamState, event: SSEEvent) -> None: """Accumulate delta content based on category.""" content = event.content or "" if not content: return cat = event.category or "" if cat == "tool_call_response": state.tool_call_response_content += content elif cat == "output_conclusion": if state.output_conclusion_chunks is not None: state.output_conclusion_chunks.append(content) # llm / think / llm_reasoning deltas: silent in summary mode # (they are internal reasoning, not shown in CLI) def _handle_data(state: StreamState, event: SSEEvent) -> None: """Second-level dispatch by category for data events.""" content = event.content or "" cat = event.category or "" if cat in _SKIP_DATA_CATEGORIES or not content: return if cat == "ask_plan": _handle_ask_plan(state, content) elif cat == "plan": _handle_data_plan(state, event) elif cat == "task_finish": _handle_data_task_finish(state, content) elif cat == "jsx_report": _handle_data_jsx_report(state, content) elif cat == "mission_report": _handle_data_mission_report(state, content) elif cat == "recommended_question": out = _fmt_recommended_questions(content) if out: _out(state, out) elif cat == "ask_report_render": out = _fmt_ask_report_render(content, getattr(state, "session_id", None)) if out: _out(state, out) state.got_content = True state.need_user_confirm = True return elif cat == "output_conclusion": # Dual-path: if content_start was received, accumulate; otherwise direct if state.output_conclusion_chunks is not None: state.output_conclusion_chunks.append(content) else: # Old format without content_start/content_finish lifecycle header = _consume_step_label(state) _out(state, _fmt_output_conclusion(content, output_dir=state.output_dir, header=header)) state.got_content = True elif cat == "tool_call_response": # data event inside content lifecycle — handled at content_finish pass else: # Fallback: extract title/step from JSON for [Step] display _handle_data_fallback(state, content) def _handle_data_plan(state: StreamState, event: SSEEvent) -> None: """Handle data/plan events — show step progress. detail mode: full "[Plan] Step n/N: name" with description. summary mode: compact progress dot on step change. """ content = event.content or "" ct = getattr(event, "content_type", None) if ct and ct != "json": return try: data = json.loads(content) if isinstance(content, str) else content except (json.JSONDecodeError, TypeError): return new_step = data.get("current_step") new_status = data.get("plan_status") # Count total steps from plan data and update step_names plans = data.get("plans", []) if plans and isinstance(plans[0], dict): inner_plan = plans[0].get("plan", {}) if isinstance(inner_plan, dict): steps = inner_plan.get("steps", []) if steps: state.total_steps = len(steps) for idx, s in enumerate(steps, 1): order = s.get("order", idx) name = s.get("name", "") if name: state.step_names[order] = name # Only react when step changes if new_step is not None and (new_step != state.current_step or new_status != state.current_plan_status): state.current_step = new_step state.current_plan_status = new_status if state.output_mode == "detail": out = _fmt_plan_progress(data, new_step, state.total_steps) if out: _log(state, out, flush=True) state.got_content = True else: # summary: build step label for next conclusion header step_name = state.step_names.get(new_step, "") total = state.total_steps progress = f"Step {new_step}/{total}" if total else f"Step {new_step}" if step_name: state.pending_step_label = f"{progress}: {step_name}" else: state.pending_step_label = progress def _handle_data_task_finish(state: StreamState, content: str) -> None: """Handle data/task_finish — structured query result.""" json_parts = _extract_json_objects(content) for _, _, parsed in json_parts: if isinstance(parsed, list) and parsed and isinstance(parsed[0], dict) and "title" in parsed[0]: out = _fmt_insights(parsed) if out: _out(state, out) state.got_content = True return elif isinstance(parsed, dict): out = _fmt_task_finish(parsed) if out: _out(state, out) state.got_content = True return def _handle_data_jsx_report(state: StreamState, content: str) -> None: """Handle data/jsx_report events. Shown in all modes.""" try: data = json.loads(content) if isinstance(content, str) else content report_type = data.get("type", "") if isinstance(data, dict) else "" _out(state, f"\n### Report Generated" + (f" ({report_type})" if report_type else "")) except (json.JSONDecodeError, TypeError): _out(state, "\n### Report Generated") def _handle_data_mission_report(state: StreamState, content: str) -> None: """Handle data/mission_report events. Shown in all modes.""" try: data = json.loads(content) if isinstance(content, str) else content title = data.get("title", "") if isinstance(data, dict) else "" _out(state, f"\n### Task Report Generated" + (f": {title}" if title else "")) except (json.JSONDecodeError, TypeError): _out(state, "\n### Task Report Generated") def _handle_data_fallback(state: StreamState, content: str) -> None: """Fallback for unrecognised data categories — extract title/step. Only shown in detail mode. """ if state.output_mode != "detail": return json_parts = _extract_json_objects(content) for _, _, parsed in json_parts: if isinstance(parsed, dict): if parsed.get("result_type") == "jupyter_cell": out = _fmt_jupyter_cell(parsed) if out: _log(state, out, flush=True) state.got_content = True elif parsed.get("title"): _log(state, f" [Step] {parsed['title']}", flush=True) state.got_content = True elif parsed.get("step"): _log(state, f" [Step] {parsed['step']}", flush=True) state.got_content = True elif parsed.get("status") and parsed.get("message"): _log(state, f" [{parsed['status']}] {parsed['message']}", flush=True) state.got_content = True def _handle_content_finish(state: StreamState, event: SSEEvent) -> None: """Flush accumulators when content lifecycle ends.""" last_category = state.content_category state.content_category = None if last_category == "tool_call_response" and state.tool_call_response_content: _flush_tool_call_response(state) elif last_category == "output_conclusion" and state.output_conclusion_chunks is not None: text = "".join(state.output_conclusion_chunks) if text.strip(): header = _consume_step_label(state) _out(state, _fmt_output_conclusion(text, output_dir=state.output_dir, header=header)) state.got_content = True state.output_conclusion_chunks = None def _flush_tool_call_response(state: StreamState) -> None: """Parse accumulated tool_call_response and format output. jupyter_cell results only shown in detail mode. """ raw = state.tool_call_response_content state.tool_call_response_content = "" try: data = json.loads(raw) except (json.JSONDecodeError, TypeError): return result_type = data.get("result_type", "") if result_type == "jupyter_cell": if state.output_mode != "detail": return # Re-parse the nested result field for jupyter cell formatting result_raw = data.get("result", "") try: inner = json.loads(result_raw) if isinstance(result_raw, str) else result_raw except (json.JSONDecodeError, TypeError): inner = result_raw if isinstance(inner, dict): out = _fmt_jupyter_cell({"result": json.dumps(inner) if not isinstance(result_raw, str) else result_raw, **{k: v for k, v in data.items() if k != "result"}}) if out: _log(state, out, flush=True) state.got_content = True elif result_type == "plan": # Plan results from tool_call_response — parse and display result_raw = data.get("result", "") try: plan_data = json.loads(result_raw) if isinstance(result_raw, str) else result_raw except (json.JSONDecodeError, TypeError): return if isinstance(plan_data, dict): _handle_data_plan(state, SSEEvent( event_type="data", data={}, category="plan", content=json.dumps(plan_data), content_type="json", )) # Plan generation in tool call needs user confirmation if not state.is_attach: _out(state, f"\n> ⚠️ Please review the execution plan above. To confirm, DO NOT create a new session, use the existing session:") session_id = state.session_id or "" _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '确认执行'") _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q 'confirm'") state.got_content = True state.need_user_confirm = True else: # When in attach mode (plan confirmed), continue processing without setting need_user_confirm _out(state, f"\n> ⚠️ Plan confirmed, continuing analysis...") state.got_content = True elif result_type == "empty": pass # ignore # Other result types: silent (no useful CLI output) def _handle_status_change(state: StreamState, event: SSEEvent) -> None: """Handle status_change events. Only visible in detail mode.""" content = event.content or "" ct = getattr(event, "content_type", None) if ct and ct != "json": return # Always track state internally try: data = json.loads(content) if isinstance(content, str) else content if isinstance(data, dict): state.current_task = data.get("current_task") except (json.JSONDecodeError, TypeError): pass # Only print in detail mode if state.output_mode == "detail": out = _fmt_status_change(content) if out: _log(state, out, flush=True) def _handle_chat_finish(state: StreamState, event: SSEEvent) -> None: """Handle chat_finish events — completions and user confirmations.""" content = event.content or "" cat = event.category or "" if cat == "chat": q = state.current_question is_confirmation = False if q: q_lower = q.strip().lower() exact_confirms = { "确认", "confirm", "execute", "同意", "approve", "yes", "y", "确认执行", "确认执行当前sql", "同意后续所有sql执行" } if q_lower in exact_confirms: is_confirmation = True if q and not is_confirmation: # Truncate long questions label = q if len(q) <= 40 else q[:37] + "..." _out(state, f"\n✅ 「{label}」已完成分析，正在获取生成的报告和中间文件...") else: _out(state, "\n✅ 分析已完成，正在获取生成的报告和中间文件...") state.got_content = True return if cat == "ask_sql" and content: _handle_ask_sql(state, content) return if cat == "ask_plan" and content: _handle_ask_plan(state, content) # If in attach mode, ensure need_user_confirm is set to False after plan confirmation if state.is_attach: state.need_user_confirm = False return if cat == "ask_human" and content: if state.is_attach: # When in attach mode, we should continue processing further events _out(state, f"\n[Processing human input response]") _out(state, f"{content}") _out(state, f"\n\u26a0\ufe0f Continuing analysis after user input...") state.got_content = True # Do not set need_user_confirm to True as user has responded # Explicitly reset need_user_confirm since user has responded state.need_user_confirm = False return _out(state, f"\n[Human Input Required]") _out(state, f"{content}") _out(state, f"\n\u26a0\ufe0f Please respond using the existing session (DO NOT create a new session):") session_id = state.session_id or "" _out(state, f" python3 scripts/data_agent_cli.py attach --session-id {session_id} -q ''") state.got_content = True state.need_user_confirm = True return if cat == "ask_report_render" and content: # Avoid printing ask_report_render again if session status is already WAIT_INPUT, which is handled in attach code if state.is_attach: # When in attach mode, process the report render request out = _fmt_ask_report_render(content, getattr(state, "session_id", None)) if out: _out(state, out) _out(state, f"\n\u26a0\ufe0f Report render confirmed, continuing...") state.got_content = True # Do not set need_user_confirm to True as user has already responded # Explicitly reset need_user_confirm since user has responded state.need_user_confirm = False return out = _fmt_ask_report_render(content, getattr(state, "session_id", None)) if out: _out(state, out) state.got_content = True state.need_user_confirm = True return def _handle_ask_sql(state: StreamState, content: str) -> None: """Format ask_sql confirmation prompt.""" try: data = json.loads(content) if isinstance(content, str) else content if isinstance(data, dict): sql = data.get("sql", data.get("query", "")) question = data.get("question", "") explain_result = data.get("explain_result", "") if sql: _out(state, f"\n### Generated SQL") _out(state, f"```sql\n{sql}\n```") if question: _out(state, f"\n> **Warning**: {question}") if explain_result: _out(state, f"\n#### Explain Result\n{explain_result}") if not state.is_attach: session_id = state.session_id or "" _out(state, f"\n> ⚠️ Please review the SQL above.") _out(state, f"> To confirm and execute ONLY this SQL, DO NOT create a new session, use the existing session:") _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '确认执行当前SQL'") _out(state, f"> To agree to execute all subsequent SQL automatically:") _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '同意后续所有SQL执行'") state.got_content = True state.need_user_confirm = True return except (json.JSONDecodeError, TypeError): pass # Fallback: raw display _out(state, f"\n### Generated SQL") _out(state, f"```sql\n{content}\n```") if not state.is_attach: session_id = state.session_id or "" _out(state, f"\n> ⚠️ Please review the SQL above.") _out(state, f"> To confirm and execute ONLY this SQL, DO NOT create a new session, use the existing session:") _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '确认执行当前SQL'") _out(state, f"> To agree to execute all subsequent SQL automatically:") _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '同意后续所有SQL执行'") state.got_content = True state.need_user_confirm = True def _handle_ask_plan(state: StreamState, content: str) -> None: """Format ask_plan confirmation prompt.""" try: data = json.loads(content) if isinstance(content, str) else content if isinstance(data, dict): plans_data = data.get("plans", []) plan_id = data.get("plan_id", "") _out(state, f"\n### Execution Plan (ID: {plan_id[:16]}...)") for plan_item in plans_data: plan = plan_item.get("plan", {}) steps = plan.get("steps", []) if steps: state.total_steps = len(steps) for idx, step in enumerate(steps, 1): order = step.get("order", idx) name = step.get("name", "") if name: state.step_names[order] = name for step in steps: order = step.get("order", "?") name = step.get("name", "") desc = step.get("description", "") step_type = step.get("type", "") status = step.get("status", "") _out(state, f"\n#### Step {order}: {name}") if step_type: _out(state, f" Type: {step_type}") if desc: _out(state, f" Description: {desc}") if status: _out(state, f" Status: {status}") if not state.is_attach: _out(state, f"\n> \u26a0\ufe0f Please review the execution plan above. To confirm, DO NOT create a new session, use the existing session:") session_id = state.session_id or "" _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '\u786e\u8ba4\u6267\u884c'") _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q 'confirm'") state.got_content = True state.need_user_confirm = True else: # When in attach mode (plan confirmed), reset the need_user_confirm flag # so subsequent steps and results are processed normally _out(state, f"\n> \u26a0\ufe0f Plan confirmed, continuing analysis...") state.got_content = True # In attach mode, explicitly reset need_user_confirm since user has confirmed state.need_user_confirm = False return except (json.JSONDecodeError, TypeError): pass # Fallback _out(state, f"\n### Execution Plan") _out(state, f"{content[:1000]}...") if not state.is_attach: _out(state, f"\n> \u26a0\ufe0f Please review the plan above. To confirm, DO NOT create a new session, use the existing session:") session_id = state.session_id or "" _out(state, f"> python3 scripts/data_agent_cli.py attach --session-id {session_id} -q '\u786e\u8ba4\u6267\u884c'") state.got_content = True state.need_user_confirm = True else: # In attach mode, indicate that plan was confirmed and continue _out(state, f"\n> \u26a0\ufe0f Plan confirmed, continuing...") state.got_content = True # In attach mode, explicitly reset need_user_confirm since user has confirmed state.need_user_confirm = False def _handle_sse_failure(state: StreamState, event: SSEEvent) -> None: """Handle SSE_FAILURE events.""" content = event.content or "" ct = getattr(event, "content_type", None) error_msg = content if ct == "json": try: data = json.loads(content) error_msg = data.get("message", data.get("error", content)) except (json.JSONDecodeError, TypeError): pass _out(state, f"\n[Error] {error_msg}") state.got_content = True FILE:scripts/cli/streaming_utils.py """Common utilities for streaming and session management in Data Agent CLI.""" import os import sys import json from pathlib import Path from typing import Tuple, Optional, List from cli.log_handler import StructuredLogHandler from cli.worker_lock import acquire_worker_lock, release_worker_lock from cli.worker_utils import handle_worker_completion, get_worker_session_details, initialize_components from data_agent import SessionManager, MessageHandler, DataSource def execute_query_batch( message_handler: MessageHandler, session, queries: List[str], output_mode: str = "summary", output_dir: Optional[Path] = None, data_source: Optional[DataSource] = None, process_log_handler: Optional[StructuredLogHandler] = None ) -> tuple[bool, bool, str]: """Execute a batch of queries with streaming output. Args: message_handler: Message handler instance session: Session object queries: List of queries to execute output_mode: Output mode (summary/detail/raw) output_dir: Output directory for files data_source: Optional data source for the queries process_log_handler: Optional process log handler Returns: Tuple of (got_content, need_confirm, full_text) """ got_content, need_confirm = False, False full_text = "" for i, query in enumerate(queries): if len(queries) > 1: # Only print separators for multiple queries print(f"\n{'=' * 60}") print(f"Query {i+1}/{len(queries)}: {query}") print("=" * 60) else: print(f"\nQuery: {query}") c, nc, t = _stream_response_with_data_source( message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir, process_log_handler=process_log_handler ) if c: got_content = True if t: full_text += f"\n### Query: {query}\n" + t + "\n" if nc: need_confirm = True break # Stop processing if confirmation is needed return got_content, need_confirm, full_text def _stream_response_with_data_source( message_handler: MessageHandler, session, query: str, data_source: Optional[DataSource] = None, output_mode: str = "summary", output_dir: Optional[Path] = None, process_log_handler: Optional[StructuredLogHandler] = None ) -> tuple[bool, bool, str]: """Wrapper for _stream_response that handles data_source parameter properly.""" # Import the streaming function from the existing module from cli.streaming import _stream_response return _stream_response( message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir, process_log_handler=process_log_handler ) def execute_single_query( message_handler: MessageHandler, session, query: str, output_mode: str = "summary", output_dir: Optional[Path] = None, data_source: Optional[DataSource] = None, process_log_handler: Optional[StructuredLogHandler] = None ) -> tuple[bool, bool, str]: """Execute a single query with streaming output. Args: message_handler: Message handler instance session: Session object query: Query to execute output_mode: Output mode (summary/detail/raw) output_dir: Output directory for files data_source: Optional data source for the query process_log_handler: Optional process log handler Returns: Tuple of (got_content, need_confirm, full_text) """ print(f"\nAnalyzing...\n") got_content, need_confirm, full_text = _stream_response_with_data_source( message_handler, session, query, data_source=data_source, output_mode=output_mode, output_dir=output_dir, process_log_handler=process_log_handler ) if not got_content: print("(No response received, please retry)") elif need_confirm: print("\n⚠️ 需要用户确认，程序将退出。请完成确认后使用会话ID继续对话。") return got_content, need_confirm, full_text def run_worker_with_handler( args, data_source: Optional[DataSource] = None, query_execution_func=None ) -> None: """Generic worker execution with common error handling and session management. Args: args: Command line arguments data_source: Optional data source for queries query_execution_func: Function to execute queries, receives (message_handler, session, args) """ from cli.worker_lock import acquire_worker_lock, release_worker_lock # Get session details from environment session_id, agent_id = get_worker_session_details() session_dir = Path(f"sessions/{session_id}") # Redirect stdout to progress.log for worker process log_file_path = session_dir / "progress.log" log_file = open(log_file_path, "a", encoding="utf-8") # Changed from "w" to "a" to append instead of overwrite sys.stdout = log_file sys.stderr = log_file print(f"[Worker] Session ID: {session_id}", flush=True) print(f"[Worker] Agent ID: {agent_id}", flush=True) acquire_worker_lock(session_dir) try: # Initialize components _, _, session_manager, message_handler = initialize_components() print("[Worker] Connecting to session...", flush=True) # Parent already verified the session is ready (AgentStatus=RUNNING), # so skip wait_for_running — DescribeDataAgentSession may report # SessionStatus=CREATING for a long time even when the session is # fully usable. session = session_manager.create_or_reuse( session_id=session_id, agent_id=agent_id, wait_for_running=False ) print(f"[Worker] Session connected (status={session.status.value})", flush=True) output_mode = getattr(args, "output", "summary") output_text = "" # Execute queries using provided function or default behavior if query_execution_func: got_content, need_confirm = query_execution_func( message_handler, session, args ) else: # Default behavior: handle query if provided, otherwise use presets # This is a simplified default - in practice, each command should provide its own execution function if hasattr(args, 'query') and args.query: got_content, need_confirm, output_text = execute_single_query( message_handler, session, args.query, output_mode=output_mode, output_dir=session_dir, data_source=data_source ) else: # Provide a simple default query execution - customize per command type default_queries = ["What can you help me with?"] got_content, need_confirm, output_text = execute_query_batch( message_handler, session, default_queries, output_mode=output_mode, output_dir=session_dir, data_source=data_source ) # Handle results handle_worker_completion(session_dir, need_confirm) except Exception as e: # Handle error handle_worker_completion(session_dir, error=e) print(f"Error: {e}", file=sys.stderr, flush=True) finally: release_worker_lock(session_dir) # Close structured logging files from cli.streaming import close_structured_logging close_structured_logging() # Close the redirected file log_file.close() sys.exit(0) FILE:scripts/cli/worker_lock.py """Worker process lock utilities for preventing concurrent workers per session. Uses a PID file (worker.pid) in the session directory. The parent checks the lock before spawning; the worker acquires it on start and releases on exit. Author: Tinker Created: 2026-03-11 """ import os import signal from pathlib import Path LOCK_FILENAME = "worker.pid" def _is_pid_alive(pid: int) -> bool: """Check if a process with the given PID is still running.""" try: os.kill(pid, 0) return True except (OSError, ProcessLookupError): return False def check_worker_lock(session_dir: Path) -> int | None: """Check if a worker is already running for this session. Returns the PID of the running worker, or None if no worker is active. Stale lock files (dead PID) are automatically cleaned up. """ lock_file = session_dir / LOCK_FILENAME if not lock_file.exists(): return None try: pid = int(lock_file.read_text().strip()) except (ValueError, OSError): # Corrupt lock file — remove it lock_file.unlink(missing_ok=True) return None if _is_pid_alive(pid): return pid # Stale lock — process is dead, clean up lock_file.unlink(missing_ok=True) return None def acquire_worker_lock(session_dir: Path) -> None: """Write the current process PID to the lock file.""" lock_file = session_dir / LOCK_FILENAME lock_file.write_text(str(os.getpid())) def release_worker_lock(session_dir: Path) -> None: """Remove the lock file if it belongs to the current process.""" lock_file = session_dir / LOCK_FILENAME try: if lock_file.exists(): pid = int(lock_file.read_text().strip()) if pid == os.getpid(): lock_file.unlink(missing_ok=True) except (ValueError, OSError): lock_file.unlink(missing_ok=True) def write_worker_pid(session_dir: Path, pid: int) -> None: """Write the spawned worker PID to the lock file (called by parent).""" lock_file = session_dir / LOCK_FILENAME lock_file.write_text(str(pid)) FILE:scripts/cli/worker_utils.py """Common utilities for async worker management in Data Agent CLI.""" import os import sys import json import subprocess import tempfile from pathlib import Path from typing import Dict, Any, Optional from cli.worker_lock import write_worker_pid, check_worker_lock from data_agent import ( DataAgentConfig, DataAgentClient, SessionManager, MessageHandler, ) def is_worker_process() -> bool: """Check if this is a worker process.""" return os.environ.get("DATA_AGENT_ASYNC_WORKER") == "1" def get_worker_session_details() -> tuple[str, str]: """Get session and agent IDs from environment variables in worker process.""" session_id = os.environ["DATA_AGENT_SESSION_ID"] agent_id = os.environ.get("DATA_AGENT_AGENT_ID", "") return session_id, agent_id def initialize_components() -> tuple: """Initialize common Data Agent components.""" config = DataAgentConfig.from_env() client = DataAgentClient(config) session_manager = SessionManager(client) message_handler = MessageHandler(client) return config, client, session_manager, message_handler def setup_async_worker( args: object, session, additional_env_vars: Optional[Dict[str, str]] = None ) -> None: """Setup and run async worker process. Args: args: Command line arguments session: Session object containing session_id and agent_id additional_env_vars: Additional environment variables to set """ session_id = session.session_id session_dir = Path(f"sessions/{session_id}") session_dir.mkdir(parents=True, exist_ok=True) # Check for existing worker existing_pid = check_worker_lock(session_dir) if existing_pid: print(f"⚠️ A worker process (PID {existing_pid}) is already running for session {session_id}.", file=sys.stderr) print(f" Check progress: cat sessions/{session_id}/progress.log", file=sys.stderr) print(f" Current status: {(session_dir / 'status.txt').read_text().strip() if (session_dir / 'status.txt').exists() else 'unknown'}", file=sys.stderr) sys.exit(1) # Save input.json input_data = vars(args).copy() input_data.pop("func", None) with open(session_dir / "input.json", "w", encoding="utf-8") as f: json.dump(input_data, f, ensure_ascii=False, indent=2) # Write status.txt with open(session_dir / "status.txt", "w", encoding="utf-8") as f: f.write("running") # Construct worker command cmd = [sys.executable] + [arg for arg in sys.argv if arg != "--async-run"] env = os.environ.copy() env["DATA_AGENT_ASYNC_WORKER"] = "1" env["DATA_AGENT_SESSION_ID"] = session_id env["DATA_AGENT_AGENT_ID"] = session.agent_id # Set additional environment variables if provided if additional_env_vars: env.update(additional_env_vars) # Make sure PYTHONIOENCODING is set so unicode is flushed properly env["PYTHONIOENCODING"] = "utf-8" # Force unbuffered output so logs appear immediately env["PYTHONUNBUFFERED"] = "1" # Start the subprocess in background and immediately write the PID # Worker will redirect its own stdout to progress.log proc = subprocess.Popen( cmd, stdout=subprocess.DEVNULL, # Worker handles its own logging stderr=subprocess.DEVNULL, start_new_session=True, env=env ) # Write PID immediately after starting the process write_worker_pid(session_dir, proc.pid) # For true async behavior, we return immediately without waiting # The actual logging will be handled by the worker process itself print(f"\n✅ Async task started. Session ID: {session_id}") print(f"Check progress at: sessions/{session_id}/progress.log") sys.exit(0) def handle_worker_completion(session_dir: Path, need_confirm: bool = False, error: Optional[Exception] = None) -> None: """Handle worker process completion by updating status files. Args: session_dir: Session directory path need_confirm: Whether user confirmation is needed error: Exception object if there was an error, None otherwise """ if error: # Error occurred with open(session_dir / "status.txt", "w", encoding="utf-8") as f: f.write("failed") with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({"status": "failed", "error": str(error)}, f) elif need_confirm: # Need user confirmation with open(session_dir / "status.txt", "w", encoding="utf-8") as f: f.write("waiting_input") with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({"status": "waiting_input"}, f) else: # Success with open(session_dir / "status.txt", "w", encoding="utf-8") as f: f.write("completed") with open(session_dir / "result.json", "w", encoding="utf-8") as f: json.dump({"status": "completed"}, f) FILE:scripts/data_agent/__init__.py """ Alibaba Cloud Yaoichi Data Agent Python SDK A Python SDK for interacting with Alibaba Cloud's Data Agent service, enabling natural language data analysis across various data sources. Author: Tinker Created: 2026-03-01 """ from data_agent.config import DataAgentConfig from data_agent.client import DataAgentClient, AsyncDataAgentClient from data_agent.session import SessionManager, AsyncSessionManager from data_agent.message import MessageHandler, AsyncMessageHandler from data_agent.file_manager import FileManager, AsyncFileManager from data_agent.sse_client import SSEClient, AsyncSSEClient, SSEEvent from data_agent.models import SessionInfo, ContentBlock, FileInfo, DataSource from data_agent.mcp_tools import DmsMcpTools, AsyncDmsMcpTools, DmsInstance, AskDatabaseResult, PagedResult from data_agent.api_adapter import APIAdapter from data_agent.exceptions import ( DataAgentException, ConfigurationError, AuthenticationError, SessionError, SessionTimeoutError, MessageError, ContentFetchError, FileUploadError, ApiError, ) __version__ = "0.1.0" __all__ = [ # Config "DataAgentConfig", # Clients "DataAgentClient", "AsyncDataAgentClient", # SSE Clients "SSEClient", "AsyncSSEClient", "SSEEvent", # Session Management "SessionManager", "AsyncSessionManager", # Message Handling "MessageHandler", "AsyncMessageHandler", # File Management "FileManager", "AsyncFileManager", # MCP Tools "DmsMcpTools", "AsyncDmsMcpTools", "PagedResult", "DmsInstance", "AskDatabaseResult", # API Adapter "APIAdapter", # Models "SessionInfo", "ContentBlock", "FileInfo", "DataSource", # Exceptions "DataAgentException", "ConfigurationError", "AuthenticationError", "SessionError", "SessionTimeoutError", "MessageError", "ContentFetchError", "FileUploadError", "ApiError", ] FILE:scripts/data_agent/api_adapter.py """API adapter for consistent casing in requests and responses. Author: Tinker Created: 2026-03-19 """ import json from typing import Any, Dict, Union def camel_to_pascal_case(s: str) -> str: """Convert camelCase string to PascalCase.""" if not s: return s return s[0].upper() + s[1:] if len(s) > 1 else s.upper() def pascal_to_camel_case(s: str) -> str: """Convert PascalCase string to camelCase.""" if not s: return s return s[0].lower() + s[1:] if len(s) > 1 else s.lower() def convert_keys_to_pascal(obj: Any, exclude_paths: list = None) -> Any: """Recursively convert all dictionary keys to PascalCase. Args: obj: Input object (dict, list, or primitive) exclude_paths: List of key names to exclude from conversion Returns: Object with keys converted to PascalCase """ if exclude_paths is None: exclude_paths = [] if isinstance(obj, dict): result = {} for key, value in obj.items(): if isinstance(key, str) and key in exclude_paths: # Skip conversion for excluded keys pascal_key = key else: pascal_key = camel_to_pascal_case(key) if isinstance(key, str) else key result[pascal_key] = convert_keys_to_pascal(value, exclude_paths) return result elif isinstance(obj, list): return [convert_keys_to_pascal(item, exclude_paths) for item in obj] else: return obj def convert_keys_to_camel(obj: Any, api_action: str = None) -> Any: """Recursively convert all dictionary keys to camelCase. Args: obj: Input object (dict, list, or primitive) api_action: API action name to customize transformation logic Returns: Object with keys converted to camelCase """ # Special handling for file upload signature response which needs to preserve specific formats if api_action and 'DescribeFileUploadSignature' in api_action: # Don't transform the response for file upload signature as it needs specific format return obj if isinstance(obj, dict): result = {} for key, value in obj.items(): camel_key = pascal_to_camel_case(key) if isinstance(key, str) else key result[camel_key] = convert_keys_to_camel(value, api_action) return result elif isinstance(obj, list): return [convert_keys_to_camel(item, api_action) for item in obj] else: return obj class APIAdapter: """Adapter to ensure consistent casing for API requests and responses.""" @staticmethod def prepare_request_params(params: Dict[str, Any], api_action: str = None) -> Dict[str, Any]: """Prepare request parameters with PascalCase keys. Args: params: Original request parameters api_action: API action name to customize transformation logic Returns: Parameters with PascalCase keys """ # Special handling for file upload signatures which may require specific param formats exclude_paths = [] if api_action and 'FileUpload' in api_action: # Some file upload APIs might expect specific parameter casing exclude_paths = ['FileName', 'FileSize', 'FileId', 'Filename'] return convert_keys_to_pascal(params, exclude_paths) @staticmethod def prepare_request_body(body: Dict[str, Any]) -> Dict[str, Any]: """Prepare request body with PascalCase keys. Args: body: Original request body Returns: Body with PascalCase keys """ return convert_keys_to_pascal(body) @staticmethod def process_response(response: Dict[str, Any], api_action: str = None) -> Dict[str, Any]: """Process API response to convert keys to camelCase. Args: response: Original API response api_action: API action name to customize transformation logic Returns: Response with camelCase keys """ return convert_keys_to_camel(response, api_action) # Singleton instance adapter = APIAdapter() FILE:scripts/data_agent/client.py """Data Agent client for API interactions. Author: Tinker Created: 2026-03-01 """ from __future__ import annotations import asyncio import time import functools import json import os from typing import Optional, Any, Callable, TypeVar import requests from alibabacloud_tea_openapi import models as open_api_models from alibabacloud_tea_openapi.client import Client as OpenApiClient from alibabacloud_tea_util import models as util_models from alibabacloud_openapi_util.client import Client as OpenApiUtilClient from Tea.exceptions import TeaException from data_agent.config import DataAgentConfig from data_agent.models import SessionInfo, SessionStatus, DataSource from data_agent.api_adapter import APIAdapter from data_agent.exceptions import ( ApiError, AuthenticationError, SessionCreationError, ConfigurationError, ) T = TypeVar("T") def retry_on_error(max_retries: int = 3, retry_codes: tuple = ("Throttling", "ServiceUnavailable")): """Decorator to retry API calls on transient errors with exponential backoff.""" def decorator(func: Callable[..., T]) -> Callable[..., T]: @functools.wraps(func) def wrapper(self, *args, **kwargs) -> T: last_exception = None for attempt in range(max_retries + 1): try: return func(self, *args, **kwargs) except ApiError as e: last_exception = e if e.code not in retry_codes or attempt == max_retries: raise wait_time = 2**attempt time.sleep(wait_time) raise last_exception return wrapper return decorator class DataAgentClient: """Synchronous client for Data Agent API. This client wraps the Alibaba Cloud DMS SDK and provides methods to interact with Data Agent sessions. """ def __init__(self, config: DataAgentConfig): """Initialize the Data Agent client. Args: config: Configuration for the client. """ self._config = config self._sdk_client: Optional[OpenApiClient] = None self._initialize_client() def _initialize_client(self) -> None: """Initialize the underlying SDK client based on authentication type.""" # Determine authentication type if self._config.api_key: # API_KEY authentication - don't initialize the Tea SDK client self._sdk_client = None self._auth_type = "api_key" else: # Use Alibaba Cloud default credential chain # Supports: env vars, ~/.aliyun/config.json, ECS role, OIDC role, etc. from alibabacloud_credentials.client import Client as CredentialClient try: self._credential_client = CredentialClient() credential = self._credential_client.get_credential() sdk_config = open_api_models.Config() sdk_config.endpoint = self._config.endpoint sdk_config.user_agent = "AlibabaCloud-Agent-Skills" sdk_config.credential = self._credential_client self._sdk_client = OpenApiClient(sdk_config) self._auth_type = "default_credential_chain" except Exception as e: raise AuthenticationError( f"Failed to get credentials from default credential chain: {e}. " "Please configure credentials via ~/.aliyun/config.json, " "environment variables, or instance role." ) def _call_api( self, action: str, version: str, params: dict, method: str = "POST", body: dict = None, ) -> dict: """Make a generic API call. Args: action: API action name. version: API version. params: Request parameters (for query string). method: HTTP method (default: POST). body: Request body (for JSON body). Returns: API response as dictionary. Raises: ApiError: If the API call fails. AuthenticationError: If authentication fails. """ if self._auth_type == "api_key": # Handle API_KEY authentication using direct HTTP requests return self._call_api_with_api_key(action, version, params, method, body) else: # Handle traditional AK/SK authentication using Tea SDK return self._call_api_with_ak_sk(action, version, params, method, body) def _call_api_with_api_key( self, action: str, version: str, params: dict, method: str = "POST", body: dict = None, ) -> dict: """Make an API call using API_KEY authentication.""" # Determine if this is a control plane or data plane API based on action control_plane_actions = [ 'ListDataAgentSession', 'CreateDataAgentSession', 'DescribeDataAgentSession', 'DescribeFileUploadSignature', 'FileUploadCallback' ] data_plane_actions = [ 'SendChatMessage', 'GetChatContent', 'DescribeDataAgentUsage', 'UpdateDataAgentSession', 'ListFileUpload', 'CreateDataAgentFeedback' ] # Choose the correct endpoint based on the action type if action in control_plane_actions: # Use control plane endpoint format - FIX: hyphen instead of dot base_endpoint = f"dataagent-{self._config.region}.aliyuncs.com/apikey" elif action in data_plane_actions: # Use data plane endpoint format - FIX: hyphen instead of dot base_endpoint = f"dataagent-stream-{self._config.region}.aliyuncs.com/apikey" else: # Default to control plane if action is not recognized - FIX: hyphen instead of dot base_endpoint = f"dataagent-{self._config.region}.aliyuncs.com/apikey" # Add the required RegionId parameter to params params['RegionId'] = self._config.region # Prepare request with PascalCase parameters prepared_params = APIAdapter.prepare_request_params(params, api_action=action) # For API_KEY authentication, we should put Action and Version in the body # according to standard OpenAPI practices if body is None: body = {} # Add action and version to the body instead of query params body.update({ 'Action': action, 'Version': version }) # Update body with prepared parameters body.update(prepared_params) prepared_body = APIAdapter.prepare_request_body(body) # Build the URL for the API call (without Action and Version in query) base_url = f"https://{base_endpoint}" # Construct the URL without action/version in query string since they're in body full_url = base_url # Set up headers with API_KEY headers = { 'x-api-key': self._config.api_key, 'Content-Type': 'application/json', 'Accept': 'application/json', 'User-Agent': 'AlibabaCloud-Agent-Skills', } # Add debug logging if enabled via environment variable if os.getenv("DATA_AGENT_DEBUG_API", "").lower() in ('true', '1', 'yes'): import pprint auth_type_display = "API_KEY" if self._auth_type == "api_key" else "AK/SK" print(f"[DEBUG] API Call: {action}") print(f"[DEBUG] Authentication Type: {auth_type_display}") print(f"[DEBUG] Method: {method}") print(f"[DEBUG] URL: {full_url}") print(f"[DEBUG] Headers: {pprint.pformat(headers)}") if prepared_body: print(f"[DEBUG] Body: {pprint.pformat(prepared_body)}") try: # Make the API call - Action and Version are now in the body if method.upper() == "GET": # For GET requests, we might still need to put action/version in query query_params = {'Action': action, 'Version': version} import urllib.parse query_string = urllib.parse.urlencode(query_params) full_url_with_params = f"{full_url}?{query_string}" response = requests.get(full_url_with_params, headers=headers, timeout=self._config.timeout) elif method.upper() == "POST": # For POST requests, Action and Version are in the body response = requests.post(full_url, headers=headers, json=prepared_body, timeout=self._config.timeout) else: # For other methods, default to POST with body response = requests.request(method, full_url, headers=headers, json=prepared_body, timeout=self._config.timeout) # Check response status response.raise_for_status() # Parse response JSON response_data = response.json() # Process response to convert keys to camelCase processed_response = APIAdapter.process_response(response_data, api_action=action) # Add debug logging for response if enabled if os.getenv("DATA_AGENT_DEBUG_API", "").lower() in ('true', '1', 'yes'): import pprint print(f"[DEBUG] Response for {action}: {pprint.pformat(processed_response)}") return processed_response except requests.exceptions.RequestException as e: # Handle HTTP request errors error_msg = f"API call failed: {str(e)}" if hasattr(e, 'response') and e.response is not None: error_msg += f", Status: {e.response.status_code}, Body: {e.response.text[:500]}" raise ApiError(error_msg, code="HTTPRequestError", request_id=None) def _call_api_with_ak_sk( self, action: str, version: str, params: dict, method: str = "POST", body: dict = None, ) -> dict: """Make an API call using AK/SK authentication.""" # Prepare request with PascalCase parameters prepared_params = APIAdapter.prepare_request_params(params, api_action=action) prepared_body = APIAdapter.prepare_request_body(body) if body else None api_params = open_api_models.Params( action=action, version=version, protocol="HTTPS", method=method, auth_type="AK", style="RPC", pathname="/", req_body_type="json", body_type="json", ) request = open_api_models.OpenApiRequest( query=OpenApiUtilClient.query(prepared_params), body=prepared_body, ) runtime = util_models.RuntimeOptions( read_timeout=self._config.timeout * 1000, connect_timeout=30000, ) # Add debug logging if enabled via environment variable if os.getenv("DATA_AGENT_DEBUG_API", "").lower() in ('true', '1', 'yes'): import pprint auth_type_display = "API_KEY" if self._auth_type == "api_key" else "AK/SK" print(f"[DEBUG] API Call: {action}") print(f"[DEBUG] Authentication Type: {auth_type_display}") print(f"[DEBUG] Method: {method}") print(f"[DEBUG] Params: {pprint.pformat(prepared_params)}") if prepared_body: print(f"[DEBUG] Body: {pprint.pformat(prepared_body)}") try: response = self._sdk_client.call_api(api_params, request, runtime) # Process response to convert keys to camelCase processed_response = APIAdapter.process_response(response.get("body", {}), api_action=action) # Add debug logging for response if enabled if os.getenv("DATA_AGENT_DEBUG_API", "").lower() in ('true', '1', 'yes'): import pprint print(f"[DEBUG] Response for {action}: {pprint.pformat(processed_response)}") return processed_response except TeaException as e: self._handle_tea_exception(e) def _handle_tea_exception(self, e: TeaException) -> None: """Convert TeaException to appropriate custom exception. Args: e: The TeaException to handle. Raises: AuthenticationError: For authentication-related errors. ApiError: For other API errors. """ code = getattr(e, "code", "Unknown") message = getattr(e, "message", str(e)) request_id = None if hasattr(e, "data") and isinstance(e.data, dict): request_id = e.data.get("RequestId") auth_error_codes = ("InvalidAccessKeyId.NotFound", "SignatureDoesNotMatch", "Forbidden") if code in auth_error_codes: raise AuthenticationError(message, code=code, request_id=request_id) raise ApiError(message, code=code, request_id=request_id) @retry_on_error(max_retries=3) def create_session( self, database_id: Optional[str] = None, title: str = "data-agent-session", mode: Optional[str] = None, enable_search: bool = False, file_id: Optional[str] = None, # 添加文件ID参数 ) -> SessionInfo: """Create a new Data Agent session. Args: database_id: Optional database ID to bind to the session. title: Session title (required by API). mode: Optional session mode, such as "ASK_DATA", "ANALYSIS", "INSIGHT". enable_search: Whether to enable search capability in the session. file_id: Optional file ID for file-based analysis session. Returns: SessionInfo with agent_id and session_id. Raises: SessionCreationError: If session creation fails. """ params = { "Title": title, "DMSUnit": self._config.region, } if database_id: params["DatabaseId"] = database_id if mode: params["Mode"] = mode if file_id: # 当指定了文件ID时，会话将是基于文件的分析 params["File"] = file_id # Ensure session configuration (language/mode/search) is persisted on server session_config = {"Language": "CHINESE", "EnableSearch": enable_search} if mode: session_config["Mode"] = mode # For API_KEY auth, SessionConfig should be a JSON object, not string # For AK/SK auth, SessionConfig should be a JSON string if self._auth_type == "api_key": params["SessionConfig"] = session_config else: params["SessionConfig"] = json.dumps(session_config) try: response = self._call_api( action="CreateDataAgentSession", version="2025-04-14", params=params, ) # Response data is nested under 'Data' field # After API adapter processing, the keys are in camelCase data = response.get("data", response) # Changed from "Data" to "data" agent_id = data.get("agentId", "") # Changed from "AgentId" to "agentId" session_id = data.get("sessionId", "") # Changed from "SessionId" to "sessionId" agent_status = data.get("agentStatus", "").upper() # Changed from "AgentStatus" to "agentStatus" if not agent_id or not session_id: raise SessionCreationError( f"Invalid response: missing AgentId or SessionId. Response: {response}" ) # Map AgentStatus to SessionStatus. # AgentStatus reflects the underlying agent compute readiness: # when RUNNING, the session can accept messages even though # DescribeDataAgentSession may still report SessionStatus as # CREATING for a prolonged period. if agent_status == "RUNNING": status = SessionStatus.RUNNING elif agent_status == "STOPPED": status = SessionStatus.STOPPED elif agent_status == "FAILED": status = SessionStatus.FAILED else: status = SessionStatus.CREATING return SessionInfo( agent_id=agent_id, session_id=session_id, status=status, database_id=database_id, ) except ApiError as e: raise SessionCreationError(f"Failed to create session: {e.message}", request_id=e.request_id) @retry_on_error(max_retries=3) def describe_session(self, session_id: str, agent_id: str = "") -> SessionInfo: """Get the status of a session. Args: session_id: The session ID. agent_id: The agent ID (optional, but used to construct the request properly). Returns: SessionInfo with current status. """ params = { "SessionId": session_id, "DMSUnit": self._config.region, } # Only include AgentId in the request if it's provided # According to API spec, DescribeDataAgentSession doesn't require AgentId if agent_id: params["AgentId"] = agent_id response = self._call_api( action="DescribeDataAgentSession", version="2025-04-14", params=params, ) request_id = response.get("requestId", "") # Changed from "RequestId" to "requestId" data = response.get("data", response) # Changed from "Data" to "data" status_str = data.get("sessionStatus", data.get("status", "CREATING")) # Changed from "SessionStatus"/"Status" to "sessionStatus"/"status" try: status = SessionStatus(status_str) except ValueError: status = SessionStatus.CREATING # Capture the real AgentId from response if available # The agent_id in the response should take precedence over the one passed in real_agent_id = data.get("agentId") or data.get("AgentId") or agent_id # Try both transformed and original return SessionInfo( agent_id=real_agent_id, session_id=session_id, status=status, database_id=data.get("databaseId"), # Changed from "DatabaseId" to "databaseId" request_id=request_id, ) @retry_on_error(max_retries=3) def send_message( self, agent_id: str, session_id: str, message: str, message_type: str = "primary", data_source: Optional[DataSource] = None, language: str = "CHINESE", ) -> dict: """Send a message to the Data Agent. Args: agent_id: The agent ID. session_id: The session ID. message: The user's natural language query. message_type: Message type (default: "primary"). data_source: Optional DataSource with database metadata. language: Response language (default: "CHINESE"). Returns: Response from the API. """ # Query parameters for RPC style API params = { "AgentId": agent_id, "SessionId": session_id, "Message": message, "MessageType": message_type, "DMSUnit": self._config.region, } # SessionConfig format depends on auth type # For API_KEY auth: JSON object # For AK/SK auth: JSON string session_config = {"Language": language} if self._auth_type == "api_key": params["SessionConfig"] = session_config else: params["SessionConfig"] = json.dumps(session_config) # DataSource format depends on auth type # For API_KEY auth: JSON object # For AK/SK auth: JSON string if data_source: if self._auth_type == "api_key": params["DataSource"] = data_source.to_api_dict() else: params["DataSource"] = json.dumps(data_source.to_api_dict()) return self._call_api( action="SendChatMessage", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def get_chat_content( self, agent_id: str, session_id: str, checkpoint: Optional[str] = None, ) -> dict: """Get chat content from the Data Agent. Args: agent_id: The agent ID. session_id: The session ID. checkpoint: Optional checkpoint for incremental fetching. Returns: Response containing content blocks. """ params = { "AgentId": agent_id, "SessionId": session_id, } if checkpoint: params["Checkpoint"] = checkpoint return self._call_api( action="GetChatContent", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def get_file_upload_signature( self, filename: str, file_size: int, ) -> dict: """Get OSS upload signature for file upload. Args: filename: Name of the file to upload. file_size: Size of the file in bytes. Returns: Response containing upload URL and credentials. """ params = { "FileName": filename, "FileSize": file_size, } return self._call_api( action="DescribeFileUploadSignature", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def file_upload_callback(self, file_id: str, filename: str, upload_location: str, file_size: int = None) -> dict: """Notify the service that file upload is complete. Args: file_id: The file ID from upload signature (uploadDir). filename: The original filename. upload_location: The full OSS path (UploadHost/UploadDir/Filename). file_size: The file size in bytes (required for API_KEY auth). Returns: Response confirming the upload. """ params = { "Filename": filename, "UploadLocation": upload_location, "FileFrom": "Skill", } if file_size: params["FileSize"] = file_size return self._call_api( action="FileUploadCallback", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def list_files(self, session_id: str, file_category: Optional[str] = None) -> dict: """List files associated with a session. Args: session_id: The session ID. file_category: Optional filter, e.g. "WebReport" for agent-generated reports, or None to list all files. Returns: Response containing file list. """ params = { "SessionId": session_id, } if file_category: params["FileCategory"] = file_category return self._call_api( action="ListFileUpload", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def list_databases( self, search_key: Optional[str] = None, page_number: int = 1, page_size: int = 50, ) -> dict: """List databases registered in DMS Data Center. Args: search_key: Optional keyword to filter by database or instance name. page_number: Page number (1-based). page_size: Number of results per page. Returns: Raw API response containing Data.List with database metadata. """ params: dict = { "PageNumber": page_number, "PageSize": page_size, } if search_key: params["SearchKey"] = search_key return self._call_api( action="ListDataCenterDatabase", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def list_tables( self, instance_name: str, database_name: str, page_number: int = 1, page_size: int = 200, ) -> dict: """List tables inside a DMS Data Center database. Args: instance_name: The InstanceName returned by list_databases. database_name: The DatabaseName returned by list_databases. page_number: Page number (1-based). page_size: Number of results per page. Returns: Raw API response containing Data.List with table metadata. """ return self._call_api( action="ListDataCenterTable", version="2025-04-14", params={ "InstanceName": instance_name, "DatabaseName": database_name, "PageNumber": page_number, "PageSize": page_size, }, ) @retry_on_error(max_retries=3) def delete_file(self, file_id: str) -> dict: """Delete an uploaded file. Args: file_id: The file ID to delete. Returns: Response confirming deletion. """ params = { "FileId": file_id, } return self._call_api( action="DeleteFileUpload", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def add_data_center_table( self, instance_name: str, database_name: str, dms_instance_id: int, dms_db_id: int, table_name_list: list[str], db_type: str = "mysql", region_id: Optional[str] = None, ) -> dict: """Add DMS database tables to Data Agent Data Center. This method imports DMS database tables into Data Agent's Data Center using the AddDataCenterTable API. Args: instance_name: RDS instance name (e.g., "rm-xxxxx"). database_name: Database name (e.g., "employees"). dms_instance_id: DMS instance ID (e.g., 1234567). dms_db_id: DMS database ID (e.g., 12345678). table_name_list: List of table names to import (required). db_type: Database type (default: "mysql"). region_id: Optional region ID (defaults to config.region). Returns: API response containing the import result. Raises: ApiError: If the API call fails. """ region = region_id or self._config.region # Convert table_name_list to JSON string for RPC API import json table_list_json = json.dumps(table_name_list, ensure_ascii=False) params = { "DMSUnit": region, "RegionId": region, "ImportType": "DMS", "InstanceName": instance_name, "DmsInstanceId": dms_instance_id, "DbType": db_type, "DatabaseName": database_name, "DmsDbId": dms_db_id, "TableNameList": table_list_json, } return self._call_api( action="AddDataCenterTable", version="2025-04-14", params=params, ) @retry_on_error(max_retries=3) def list_sessions( self, start_time: Optional[str] = None, end_time: Optional[str] = None, page_number: int = 1, page_size: int = 10, ) -> dict: """List Data Agent sessions. For API_KEY authentication, uses startTime/endTime parameters. For AK/SK authentication, uses createStartTime/createEndTime parameters. Args: start_time: Start time for filtering sessions (ISO format). end_time: End time for filtering sessions (ISO format). page_number: Page number (1-based). page_size: Number of results per page. Returns: Response containing list of sessions. """ params = { "PageNumber": page_number, "PageSize": page_size, } # Add time parameters based on authentication type if self._auth_type == "api_key": # For API_KEY auth, use the new parameter names if start_time: params["StartTime"] = start_time if end_time: params["EndTime"] = end_time else: # For AK/SK auth, use the original parameter names if start_time: params["CreateStartTime"] = start_time if end_time: params["CreateEndTime"] = end_time return self._call_api( action="ListDataAgentSession", version="2025-04-14", params=params, ) @property def config(self) -> DataAgentConfig: """Get the client configuration.""" return self._config class AsyncDataAgentClient: """Asynchronous client for Data Agent API. This client provides async/await support for all API operations. """ def __init__(self, config: DataAgentConfig): """Initialize the async Data Agent client. Args: config: Configuration for the client. """ self._config = config self._sync_client = DataAgentClient(config) async def _run_in_executor(self, func: Callable[..., T], *args, **kwargs) -> T: """Run a synchronous function in an executor. Args: func: The function to run. *args: Positional arguments. **kwargs: Keyword arguments. Returns: The result of the function. """ loop = asyncio.get_event_loop() return await loop.run_in_executor( None, functools.partial(func, *args, **kwargs), ) async def create_session( self, database_id: Optional[str] = None, title: str = "data-agent-session", mode: Optional[str] = None, enable_search: bool = False, file_id: Optional[str] = None, # 添加文件ID参数 ) -> SessionInfo: """Create a new Data Agent session asynchronously. Args: database_id: Optional database ID to bind to the session. title: Session title (required by API). mode: Optional session mode, such as "ASK_DATA", "ANALYSIS", "INSIGHT". enable_search: Whether to enable search capability in the session. file_id: Optional file ID for file-based analysis session. Returns: SessionInfo with agent_id and session_id. """ return await self._run_in_executor( self._sync_client.create_session, database_id=database_id, title=title, mode=mode, enable_search=enable_search, file_id=file_id, ) async def describe_session(self, session_id: str, agent_id: str = "") -> SessionInfo: """Get the status of a session asynchronously. Args: session_id: The session ID. agent_id: The agent ID (optional). Returns: SessionInfo with current status. """ return await self._run_in_executor( self._sync_client.describe_session, session_id=session_id, agent_id=agent_id, ) async def send_message( self, agent_id: str, session_id: str, message: str, message_type: str = "primary", data_source: Optional[DataSource] = None, language: str = "CHINESE", ) -> dict: """Send a message to the Data Agent asynchronously. Args: agent_id: The agent ID. session_id: The session ID. message: The user's natural language query. message_type: Message type (default: "primary"). data_source: Optional DataSource with database metadata. language: Response language (default: "CHINESE"). Returns: Response from the API. """ return await self._run_in_executor( self._sync_client.send_message, agent_id=agent_id, session_id=session_id, message=message, message_type=message_type, data_source=data_source, language=language, ) async def get_chat_content( self, agent_id: str, session_id: str, checkpoint: Optional[str] = None, ) -> dict: """Get chat content from the Data Agent asynchronously. Args: agent_id: The agent ID. session_id: The session ID. checkpoint: Optional checkpoint for incremental fetching. Returns: Response containing content blocks. """ return await self._run_in_executor( self._sync_client.get_chat_content, agent_id=agent_id, session_id=session_id, checkpoint=checkpoint, ) async def get_file_upload_signature( self, filename: str, file_size: int, ) -> dict: """Get OSS upload signature asynchronously. Args: filename: Name of the file to upload. file_size: Size of the file in bytes. Returns: Response containing upload URL and credentials. """ return await self._run_in_executor( self._sync_client.get_file_upload_signature, filename=filename, file_size=file_size, ) async def file_upload_callback(self, file_id: str, filename: str, upload_location: str, file_size: int = None) -> dict: """Notify file upload completion asynchronously. Args: file_id: The file ID from upload signature. filename: The original filename. upload_location: The full OSS path. file_size: The file size in bytes (required for API_KEY auth). Returns: Response confirming the upload. """ return await self._run_in_executor( self._sync_client.file_upload_callback, file_id=file_id, filename=filename, upload_location=upload_location, file_size=file_size, ) async def list_files(self, session_id: str, file_category: Optional[str] = None) -> dict: """List files asynchronously. Args: session_id: The session ID. file_category: Optional filter, e.g. "WebReport" for agent-generated reports. Returns: Response containing file list. """ return await self._run_in_executor( self._sync_client.list_files, session_id=session_id, file_category=file_category, ) async def delete_file(self, file_id: str) -> dict: """Delete a file asynchronously. Args: file_id: The file ID to delete. Returns: Response confirming deletion. """ return await self._run_in_executor( self._sync_client.delete_file, file_id=file_id, ) @property def config(self) -> DataAgentConfig: """Get the client configuration.""" return self._config FILE:scripts/data_agent/config.py """Configuration management for Data Agent SDK. Author: Tinker Created: 2026-03-01 """ from __future__ import annotations import os from dataclasses import dataclass, field from typing import Optional from dotenv import load_dotenv from data_agent.exceptions import ConfigurationError @dataclass class DataAgentConfig: """Configuration for Data Agent client. Attributes: api_key: API Key for alternative authentication (optional) region: Region for DMS endpoint (default: cn-hangzhou) endpoint: Custom endpoint (auto-generated if not set) timeout: API timeout in seconds (default: 300) max_retry: Maximum retry attempts (default: 3) poll_interval: Interval between polls in seconds (default: 2) max_poll_count: Maximum poll attempts (default: 60) Note: For AK/SK authentication, the SDK uses Alibaba Cloud default credential chain (environment variables, ~/.aliyun/config.json, instance role, etc.) Do NOT pass AK/SK explicitly to this config. """ api_key: Optional[str] = None region: str = "cn-hangzhou" endpoint: Optional[str] = None timeout: int = 300 max_retry: int = 3 poll_interval: int = 2 max_poll_count: int = 60 def __post_init__(self) -> None: """Generate endpoint if not provided and validate config.""" if not self.endpoint: if self.api_key: # For API key auth, use the dataagent domain format with /apikey suffix self.endpoint = f"dataagent-{self.region}.aliyuncs.com/apikey" else: # Standard DMS endpoint for AK/SK auth (uses default credential chain) self.endpoint = f"dms.{self.region}.aliyuncs.com" self.validate() def validate(self) -> None: """Validate configuration values. Raises: ConfigurationError: If required configuration is missing or invalid. """ # API key is optional - if not provided, SDK will use default credential chain pass if self.timeout <= 0: raise ConfigurationError(f"Invalid timeout value: {self.timeout}. Must be positive.") if self.max_retry < 0: raise ConfigurationError(f"Invalid max_retry value: {self.max_retry}. Must be non-negative.") if self.poll_interval <= 0: raise ConfigurationError(f"Invalid poll_interval value: {self.poll_interval}. Must be positive.") if self.max_poll_count <= 0: raise ConfigurationError(f"Invalid max_poll_count value: {self.max_poll_count}. Must be positive.") @classmethod def from_env(cls, dotenv_path: Optional[str] = None) -> DataAgentConfig: """Create configuration from environment variables. Args: dotenv_path: Optional path to .env file to load. Returns: DataAgentConfig instance. Note: AK/SK credentials are NOT read from environment variables here. The Alibaba Cloud SDK uses its default credential chain. """ if dotenv_path: load_dotenv(dotenv_path) else: load_dotenv() return cls( api_key=os.environ.get("DATA_AGENT_API_KEY") or None, region=os.environ.get("DATA_AGENT_REGION", "cn-hangzhou"), endpoint=os.environ.get("DATA_AGENT_ENDPOINT"), timeout=int(os.environ.get("DATA_AGENT_TIMEOUT", "300")), max_retry=int(os.environ.get("DATA_AGENT_MAX_RETRY", "3")), poll_interval=int(os.environ.get("DATA_AGENT_POLL_INTERVAL", "2")), max_poll_count=int(os.environ.get("DATA_AGENT_MAX_POLL_COUNT", "60")), ) @classmethod def from_dict(cls, config_dict: dict) -> DataAgentConfig: """Create configuration from a dictionary. Args: config_dict: Dictionary containing configuration values. Returns: DataAgentConfig instance. Note: Do NOT pass AK/SK credentials in config_dict. Use Alibaba Cloud default credential chain instead. """ return cls( api_key=config_dict.get("api_key"), region=config_dict.get("region", "cn-hangzhou"), endpoint=config_dict.get("endpoint"), timeout=config_dict.get("timeout", 300), max_retry=config_dict.get("max_retry", 3), poll_interval=config_dict.get("poll_interval", 2), max_poll_count=config_dict.get("max_poll_count", 60), ) def to_dict(self) -> dict: """Convert configuration to dictionary (excluding secrets). Returns: Dictionary with non-sensitive configuration values. """ result = { "region": self.region, "endpoint": self.endpoint, "timeout": self.timeout, "max_retry": self.max_retry, "poll_interval": self.poll_interval, "max_poll_count": self.max_poll_count, } # Only indicate auth type if self.api_key: result["auth_type"] = "api_key" else: result["auth_type"] = "default_credential_chain" return result def __repr__(self) -> str: """String representation (hides secrets).""" return ( f"DataAgentConfig(region='{self.region}', endpoint='{self.endpoint}', " f"timeout={self.timeout}, max_retry={self.max_retry})" ) FILE:scripts/data_agent/exceptions.py """Custom exceptions for Data Agent SDK. Author: Tinker Created: 2026-03-01 """ from typing import Optional class DataAgentException(Exception): """Base exception for all Data Agent errors.""" def __init__(self, message: str, request_id: Optional[str] = None): self.message = message self.request_id = request_id super().__init__(message) def __str__(self) -> str: if self.request_id: return f"{self.message} (RequestId: {self.request_id})" return self.message class ConfigurationError(DataAgentException): """Raised when configuration is invalid or missing.""" pass class AuthenticationError(DataAgentException): """Raised when authentication fails (invalid AccessKey, signature mismatch).""" def __init__(self, message: str, code: Optional[str] = None, request_id: Optional[str] = None): self.code = code super().__init__(message, request_id) class SessionError(DataAgentException): """Base exception for session-related errors.""" pass class SessionCreationError(SessionError): """Raised when session creation fails.""" pass class SessionTimeoutError(SessionError): """Raised when session fails to reach RUNNING state within timeout.""" def __init__(self, message: str, session_id: Optional[str] = None, waited_seconds: int = 0): self.session_id = session_id self.waited_seconds = waited_seconds super().__init__(message) class SessionNotFoundError(SessionError): """Raised when session does not exist.""" pass class MessageError(DataAgentException): """Base exception for message-related errors.""" pass class MessageSendError(MessageError): """Raised when message sending fails.""" pass class ContentFetchError(MessageError): """Raised when content fetching fails or times out.""" def __init__(self, message: str, partial_content: Optional[str] = None, request_id: Optional[str] = None): self.partial_content = partial_content super().__init__(message, request_id) class FileError(DataAgentException): """Base exception for file-related errors.""" pass class FileUploadError(FileError): """Raised when file upload fails.""" def __init__(self, message: str, file_path: Optional[str] = None, request_id: Optional[str] = None): self.file_path = file_path super().__init__(message, request_id) class FileDownloadError(FileError): """Raised when file download fails.""" pass class ApiError(DataAgentException): """Raised for general API errors with error code.""" def __init__( self, message: str, code: str, request_id: Optional[str] = None, http_status: Optional[int] = None, ): self.code = code self.http_status = http_status super().__init__(message, request_id) def __str__(self) -> str: parts = [f"[{self.code}] {self.message}"] if self.http_status: parts.append(f"(HTTP {self.http_status})") if self.request_id: parts.append(f"(RequestId: {self.request_id})") return " ".join(parts) FILE:scripts/data_agent/file_manager.py """File management for Data Agent. Author: Tinker Created: 2026-03-01 """ from __future__ import annotations import os import asyncio from pathlib import Path from typing import Optional, List import requests import aiohttp from data_agent.client import DataAgentClient, AsyncDataAgentClient from data_agent.models import FileInfo from data_agent.exceptions import FileUploadError, FileDownloadError # Supported file types and their MIME types SUPPORTED_FILE_TYPES = { ".csv": "text/csv", ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", ".xls": "application/vnd.ms-excel", ".json": "application/json", ".txt": "text/plain", } class FileManager: """Manages file uploads and downloads for Data Agent. Handles uploading local files to OSS and retrieving generated reports from Data Agent sessions. """ def __init__(self, client: DataAgentClient): """Initialize file manager. Args: client: DataAgentClient instance for API calls. """ self._client = client def upload_file( self, file_path: str, timeout: int = 60, ) -> FileInfo: """Upload a file for analysis. Args: file_path: Path to the local file. timeout: Upload timeout in seconds. Returns: FileInfo with file_id and metadata. Raises: FileUploadError: If upload fails. """ path = Path(file_path) # Validate file exists if not path.exists(): raise FileUploadError(f"File not found: {file_path}", file_path=file_path) # Validate file type suffix = path.suffix.lower() if suffix not in SUPPORTED_FILE_TYPES: raise FileUploadError( f"Unsupported file type: {suffix}. Supported types: {list(SUPPORTED_FILE_TYPES.keys())}", file_path=file_path, ) filename = path.name file_size = path.stat().st_size # Get upload signature try: signature_response = self._client.get_file_upload_signature( filename=filename, file_size=file_size, ) except Exception as e: raise FileUploadError(f"Failed to get upload signature: {e}", file_path=file_path) # Extract OSS direct upload parameters from data field (support both camelCase and PascalCase) # For API_KEY auth, the response has been processed to camelCase data = signature_response.get("data", signature_response.get("Data", {})) upload_host = data.get("uploadHost", data.get("UploadHost")) upload_dir = data.get("uploadDir", data.get("UploadDir")) policy = data.get("policy", data.get("Policy")) oss_signature = data.get("ossSignature", data.get("OssSignature")) oss_date = data.get("ossDate", data.get("OssDate")) oss_security_token = data.get("ossSecurityToken", data.get("OssSecurityToken")) oss_credential = data.get("ossCredential", data.get("OssCredential")) if not all([upload_host, upload_dir, policy, oss_signature]): raise FileUploadError( f"Invalid signature response: missing required OSS upload parameters", file_path=file_path, ) # Construct file key and upload URL file_key = f"{upload_dir}/{filename}" upload_url = upload_host # UploadLocation is the full OSS path: UploadDir/Filename upload_location = file_key # Upload to OSS using POST form (multipart/form-data) content_type = SUPPORTED_FILE_TYPES[suffix] try: with open(file_path, "rb") as f: # OSS requires file field to be named 'file' files = { "key": (None, file_key), "policy": (None, policy), "x-oss-signature": (None, oss_signature), "x-oss-date": (None, oss_date), "x-oss-security-token": (None, oss_security_token or ""), "x-oss-credential": (None, oss_credential), "x-oss-signature-version": (None, "OSS4-HMAC-SHA256"), "success_action_status": (None, "200"), "file": (filename, f, content_type), } response = requests.post( upload_url, files=files, timeout=timeout, ) response.raise_for_status() except requests.RequestException as e: raise FileUploadError(f"Failed to upload to OSS: {e}", file_path=file_path) # FileId is the UploadDir itself (the unique upload path) file_id = upload_dir # Confirm upload and get the correct FileId from response # Note: For API_KEY auth, FileUploadCallback may fail, so we use upload_dir as file_id final_file_id = file_id try: callback_resp = self._client.file_upload_callback(file_id, filename, upload_location, file_size) # Extract FileId from callback response - after API adapter transformation, # response fields become camelCase data_field = callback_resp.get("data", callback_resp.get("Data", {})) data_center_file_id = data_field.get("fileId", data_field.get("FileId")) if data_center_file_id: final_file_id = data_center_file_id except Exception as e: # If callback fails, use upload_dir as file_id (fallback) print(f"Warning: FileUploadCallback failed: {e}. Using upload_dir as file_id.") final_file_id = file_id return FileInfo( file_id=final_file_id, filename=filename, file_type=suffix.lstrip("."), size=file_size, upload_url=f"{upload_host}/{file_key}", ) def list_files(self, session_id: str, file_category: Optional[str] = None) -> List[FileInfo]: """List files associated with a session. Args: session_id: The session ID. file_category: Optional filter, e.g. "WebReport" for agent-generated reports, or None to list all files. Returns: List of FileInfo objects. """ response = self._client.list_files(session_id, file_category=file_category) files = response.get("Data", response.get("Files", [])) if not isinstance(files, list): files = [] return [ FileInfo( file_id=f.get("FileId", ""), filename=f.get("FileName", ""), file_type=f.get("FileType", ""), size=f.get("FileSize", 0), # API returns DownloadLink; fall back to DownloadUrl for compatibility download_url=f.get("DownloadLink") or f.get("DownloadUrl"), ) for f in files ] def list_reports(self, session_id: str) -> List[FileInfo]: """List agent-generated reports (WebReport) for a session. Args: session_id: The session ID. Returns: List of FileInfo objects for WebReport files. """ return self.list_files(session_id, file_category="WebReport") def download_from_url( self, download_url: str, save_path: str, timeout: int = 120, ) -> str: """Download a file from a URL and save locally. Args: download_url: Direct download URL. save_path: Local path to save the file. timeout: Download timeout in seconds. Returns: Absolute path to the saved file. Raises: FileDownloadError: If download fails. """ save = Path(save_path) save.parent.mkdir(parents=True, exist_ok=True) try: response = requests.get(download_url, timeout=timeout, stream=True) response.raise_for_status() with save.open("wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) return str(save.resolve()) except requests.RequestException as e: raise FileDownloadError(f"Failed to download file: {e}") def download_file( self, file_id: str, save_path: str, timeout: int = 60, ) -> str: """Download a file by ID. Args: file_id: The file ID to download. save_path: Path to save the downloaded file. timeout: Download timeout in seconds. Returns: Path to the saved file. Raises: FileDownloadError: If download fails. """ # Get file info with download URL # Note: This assumes the list_files API returns download URLs # Actual implementation may vary based on API behavior try: response = requests.get( f"https://dms.aliyuncs.com/files/{file_id}", # Placeholder URL timeout=timeout, stream=True, ) response.raise_for_status() with open(save_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) return save_path except requests.RequestException as e: raise FileDownloadError(f"Failed to download file: {e}") def delete_file(self, file_id: str) -> bool: """Delete an uploaded file. Args: file_id: The file ID to delete. Returns: True if deletion was successful. """ try: self._client.delete_file(file_id) return True except Exception: return False def get_file_type(self, file_path: str) -> Optional[str]: """Get the file type from path. Args: file_path: Path to the file. Returns: File type string or None if unsupported. """ suffix = Path(file_path).suffix.lower() return suffix.lstrip(".") if suffix in SUPPORTED_FILE_TYPES else None def is_supported_file(self, file_path: str) -> bool: """Check if a file type is supported. Args: file_path: Path to the file. Returns: True if file type is supported. """ suffix = Path(file_path).suffix.lower() return suffix in SUPPORTED_FILE_TYPES class AsyncFileManager: """Asynchronous file manager for Data Agent. Provides async/await support for file operations. """ def __init__(self, client: AsyncDataAgentClient): """Initialize async file manager. Args: client: AsyncDataAgentClient instance for API calls. """ self._client = client async def upload_file( self, file_path: str, timeout: int = 60, ) -> FileInfo: """Upload a file for analysis asynchronously. Args: file_path: Path to the local file. timeout: Upload timeout in seconds. Returns: FileInfo with file_id and metadata. Raises: FileUploadError: If upload fails. """ path = Path(file_path) if not path.exists(): raise FileUploadError(f"File not found: {file_path}", file_path=file_path) suffix = path.suffix.lower() if suffix not in SUPPORTED_FILE_TYPES: raise FileUploadError( f"Unsupported file type: {suffix}", file_path=file_path, ) filename = path.name file_size = path.stat().st_size # Get upload signature try: signature_response = await self._client.get_file_upload_signature( filename=filename, file_size=file_size, ) except Exception as e: raise FileUploadError(f"Failed to get upload signature: {e}", file_path=file_path) # Extract OSS direct upload parameters from data field (support both camelCase and PascalCase) # For API_KEY auth, the response has been processed to camelCase data = signature_response.get("data", signature_response.get("Data", {})) upload_host = data.get("uploadHost", data.get("UploadHost")) upload_dir = data.get("uploadDir", data.get("UploadDir")) policy = data.get("policy", data.get("Policy")) oss_signature = data.get("ossSignature", data.get("OssSignature")) oss_date = data.get("ossDate", data.get("OssDate")) oss_security_token = data.get("ossSecurityToken", data.get("OssSecurityToken")) oss_credential = data.get("ossCredential", data.get("OssCredential")) if not all([upload_host, upload_dir, policy, oss_signature]): raise FileUploadError( f"Invalid signature response: missing required OSS upload parameters", file_path=file_path, ) # Construct file key and upload URL file_key = f"{upload_dir}/{filename}" upload_url = upload_host # UploadLocation is the full OSS path: UploadDir/Filename upload_location = file_key # Upload to OSS using POST form with aiohttp content_type = SUPPORTED_FILE_TYPES[suffix] try: async with aiohttp.ClientSession() as session: with open(file_path, "rb") as f: file_content = f.read() form_data = aiohttp.FormData() form_data.add_field("key", file_key) form_data.add_field("policy", policy) form_data.add_field("x-oss-signature", oss_signature) form_data.add_field("x-oss-date", oss_date) form_data.add_field("x-oss-security-token", oss_security_token) form_data.add_field("x-oss-credential", oss_credential) form_data.add_field("x-oss-signature-version", "OSS4-HMAC-SHA256") form_data.add_field("success_action_status", "200") form_data.add_field("file", file_content, filename=filename, content_type=content_type) async with session.post( upload_url, data=form_data, timeout=aiohttp.ClientTimeout(total=timeout), ) as response: response.raise_for_status() except aiohttp.ClientError as e: raise FileUploadError(f"Failed to upload to OSS: {e}", file_path=file_path) # FileId is the UploadDir itself (the unique upload path) file_id = upload_dir # Confirm upload and get the correct FileId from response # Note: For API_KEY auth, FileUploadCallback may fail, so we use upload_dir as file_id final_file_id = file_id try: callback_resp = await self._client.file_upload_callback(file_id, filename, upload_location, file_size) # Extract FileId from callback response - after API adapter transformation, # response fields become camelCase data_field = callback_resp.get("data", callback_resp.get("Data", {})) data_center_file_id = data_field.get("fileId", data_field.get("FileId")) if data_center_file_id: final_file_id = data_center_file_id except Exception as e: # If callback fails, use upload_dir as file_id (fallback) print(f"Warning: FileUploadCallback failed: {e}. Using upload_dir as file_id.") final_file_id = file_id return FileInfo( file_id=final_file_id, filename=filename, file_type=suffix.lstrip("."), size=file_size, upload_url=upload_location, ) async def list_files(self, session_id: str) -> List[FileInfo]: """List files associated with a session asynchronously. Args: session_id: The session ID. Returns: List of FileInfo objects. """ response = await self._client.list_files(session_id) files = response.get("Files", []) return [ FileInfo( file_id=f.get("FileId", ""), filename=f.get("FileName", ""), file_type=f.get("FileType", ""), size=f.get("FileSize", 0), download_url=f.get("DownloadUrl"), ) for f in files ] async def download_file( self, download_url: str, save_path: str, timeout: int = 60, ) -> str: """Download a file from URL asynchronously. Args: download_url: URL to download from. save_path: Path to save the downloaded file. timeout: Download timeout in seconds. Returns: Path to the saved file. Raises: FileDownloadError: If download fails. """ try: async with aiohttp.ClientSession() as session: async with session.get( download_url, timeout=aiohttp.ClientTimeout(total=timeout), ) as response: response.raise_for_status() content = await response.read() with open(save_path, "wb") as f: f.write(content) return save_path except aiohttp.ClientError as e: raise FileDownloadError(f"Failed to download file: {e}") async def delete_file(self, file_id: str) -> bool: """Delete an uploaded file asynchronously. Args: file_id: The file ID to delete. Returns: True if deletion was successful. """ try: await self._client.delete_file(file_id) return True except Exception: return False FILE:scripts/data_agent/mcp_tools.py """MCP (Model Context Protocol) tools integration for DMS. This module provides integration with Alibaba Cloud DMS MCP Server tools: - listInstances: Search for database instances in DMS - askDatabase: Natural language query database (NL2SQL + execute SQL) Author: Tinker Created: 2026-03-05 """ from __future__ import annotations import os from typing import Optional, List, Dict, Any from dataclasses import dataclass from alibabacloud_dms_enterprise20181101 import models as dms_models from alibabacloud_dms_enterprise20181101.client import Client as DmsClient from alibabacloud_tea_openapi import models as open_api_models from data_agent.config import DataAgentConfig from data_agent.exceptions import ApiError @dataclass class DmsInstance: """DMS Instance information.""" instance_id: str instance_alias: str instance_type: str host: str port: int state: str env_type: str instance_source: str instance_resource_id: Optional[str] = None @dataclass class AskDatabaseResult: """Result from askDatabase query.""" sql: str result: str explanation: Optional[str] = None @dataclass class DmsDatabase: """DMS Database information.""" database_id: str schema_name: str host: str port: int instance_id: str instance_alias: str db_type: str env_type: str @dataclass class DmsTable: """DMS Table information.""" table_id: str table_name: str table_guid: str database_id: str schema_name: str engine: Optional[str] = None table_comment: Optional[str] = None @dataclass class PagedResult: """Paginated result with metadata.""" items: List[Any] page_number: int page_size: int total_count: int @property def total_pages(self) -> int: """Calculate total number of pages.""" if self.page_size <= 0: return 0 return (self.total_count + self.page_size - 1) // self.page_size @property def has_next(self) -> bool: """Check if there are more pages.""" return self.page_number < self.total_pages class DmsMcpTools: """MCP tools for DMS integration. Provides access to DMS MCP Server functionality: - listInstances: Search and list database instances - askDatabase: Natural language to SQL query execution """ def __init__(self, config: DataAgentConfig): """Initialize DMS MCP tools. Args: config: Data Agent configuration containing credentials. """ self._config = config self._dms_client: Optional[DmsClient] = None self._init_dms_client() def _init_dms_client(self) -> None: """Initialize DMS Enterprise client using default credential chain.""" from alibabacloud_credentials.client import Client as CredentialClient credential_client = CredentialClient() sdk_config = open_api_models.Config() sdk_config.credential = credential_client # Use DMS Enterprise endpoint region = self._config.region sdk_config.endpoint = f"dms-enterprise.{region}.aliyuncs.com" sdk_config.user_agent = "AlibabaCloud-Agent-Skills" self._dms_client = DmsClient(sdk_config) def list_instances( self, search_key: Optional[str] = None, db_type: Optional[str] = None, env_type: Optional[str] = None, page_number: int = 1, page_size: int = 50, ) -> PagedResult: """Search for database instances in DMS. Args: search_key: Optional search key (e.g., instance host, alias). db_type: Optional database type (e.g., mysql, polardb, oracle). env_type: Optional environment type (e.g., product, dev, test). page_number: Page number (default: 1). page_size: Page size (default: 50, max: 100). Returns: PagedResult with list of DmsInstance objects and pagination info. Raises: ApiError: If the API call fails. """ try: page_size = min(page_size, 100) # Max 100 req = dms_models.ListInstancesRequest() req.page_number = page_number req.page_size = page_size if search_key: req.search_key = search_key if db_type: req.db_type = db_type if env_type: req.env_type = env_type resp = self._dms_client.list_instances(req) instance_data = resp.body.to_map() total_count = instance_data.get("TotalCount", 0) result = [] if "InstanceList" in instance_data: instance_list = instance_data["InstanceList"] if "Instance" in instance_list: instances = instance_list["Instance"] if isinstance(instances, list): for item in instances: # Skip DELETED instances state = item.get("State", "") if state == "DELETED": continue # Handle EcsInstanceId -> InstanceResourceId mapping resource_id = item.get("EcsInstanceId") or item.get("InstanceResourceId") instance = DmsInstance( instance_id=str(item.get("InstanceId", "")), instance_alias=item.get("InstanceAlias", ""), instance_type=item.get("InstanceType", ""), host=item.get("Host", ""), port=int(item.get("Port", 0)) if item.get("Port") else 0, state=state, env_type=item.get("EnvType", ""), instance_source=item.get("InstanceSource", ""), instance_resource_id=resource_id, ) result.append(instance) return PagedResult( items=result, page_number=page_number, page_size=page_size, total_count=total_count, ) except Exception as e: raise ApiError(f"Failed to list instances: {e}", code="ListInstancesError") def ask_database( self, question: str, database_id: str, ) -> AskDatabaseResult: """Natural language query database (NL2SQL + execute SQL). Note: This is a placeholder implementation. The actual askDatabase functionality requires the Data Agent Analytics API which handles NL2SQL conversion and SQL execution. Args: question: Natural language question to query the database. database_id: DMS database ID to query. Returns: AskDatabaseResult containing the generated SQL and query result. Raises: ApiError: If the query fails. """ # This functionality is provided by the Data Agent Analytics API # which is already implemented in DataAgentClient # This method serves as a wrapper for MCP tool compatibility raise NotImplementedError( "askDatabase is implemented via DataAgentClient. " "Use DataAgentClient with ASK_DATA mode for NL2SQL queries." ) def search_database( self, search_key: str, page_number: int = 1, page_size: int = 200, ) -> List[DmsDatabase]: """Search databases by schema name in DMS. Args: search_key: Schema name search keyword (required). page_number: Page number for pagination (default: 1). page_size: Page size for pagination (default: 200, max: 1000). Returns: List of DmsDatabase objects. Raises: ApiError: If the API call fails. """ try: req = dms_models.SearchDatabaseRequest() req.search_key = search_key req.page_number = page_number req.page_size = min(page_size, 1000) # Max 1000 resp = self._dms_client.search_database(req) data = resp.body.to_map() # API returns 'SearchDatabaseList' not 'DatabaseList' db_list_key = "SearchDatabaseList" if "SearchDatabaseList" in data else "DatabaseList" if db_list_key not in data: return [] db_list = data[db_list_key] # Handle different response structures # API returns {'SearchDatabase': [...]} instead of {'Database': [...]} if isinstance(db_list, list): databases = db_list elif isinstance(db_list, dict): if "SearchDatabase" in db_list: databases = db_list["SearchDatabase"] elif "Database" in db_list: databases = db_list["Database"] else: return [] else: return [] if not isinstance(databases, list) or not databases: return [] result = [] for item in databases: # SearchDatabase API returns different fields than expected # Use Alias as instance_alias if InstanceAlias not present instance_alias = item.get("InstanceAlias", "") or item.get("Alias", "") # SearchDatabase doesn't return InstanceId, use empty string instance_id = str(item.get("InstanceId", "")) db = DmsDatabase( database_id=str(item.get("DatabaseId", "")), schema_name=item.get("SchemaName", ""), host=item.get("Host", ""), port=int(item.get("Port", 0)) if item.get("Port") else 0, instance_id=instance_id, instance_alias=instance_alias, db_type=item.get("DbType", ""), env_type=item.get("EnvType", ""), ) result.append(db) return result except Exception as e: raise ApiError(f"Failed to search database: {e}", code="SearchDatabaseError") def list_tables( self, database_id: str, search_name: Optional[str] = None, page_number: int = 1, page_size: int = 200, ) -> List[DmsTable]: """List tables in a database. Args: database_id: Database ID to search tables in (required). search_name: Table name search keyword (optional). page_number: Page number for pagination (default: 1). page_size: Page size for pagination (default: 200, max: 200). Returns: List of DmsTable objects. Raises: ApiError: If the API call fails. """ try: req = dms_models.ListTablesRequest() req.database_id = database_id if search_name: req.search_name = search_name req.page_number = page_number req.page_size = min(page_size, 200) # Max 200 resp = self._dms_client.list_tables(req) data = resp.body.to_map() if "TableList" not in data: return [] table_list = data["TableList"] if "Table" not in table_list: return [] tables = table_list["Table"] if not isinstance(tables, list) or not tables: return [] result = [] for item in tables: table = DmsTable( table_id=str(item.get("TableId", "")), table_name=item.get("TableName", ""), table_guid=item.get("TableGuid", ""), database_id=str(item.get("DatabaseId", "")), schema_name=item.get("SchemaName", ""), engine=item.get("Engine"), table_comment=item.get("TableComment"), ) result.append(table) return result except Exception as e: raise ApiError(f"Failed to list tables: {e}", code="ListTablesError") class AsyncDmsMcpTools: """Async version of DMS MCP tools.""" def __init__(self, config: DataAgentConfig): """Initialize async DMS MCP tools. Args: config: Data Agent configuration. """ self._config = config self._dms_client: Optional[DmsClient] = None async def _init_dms_client(self) -> None: """Initialize async DMS Enterprise client using default credential chain.""" from alibabacloud_credentials.client import Client as CredentialClient credential_client = CredentialClient() sdk_config = open_api_models.Config() sdk_config.credential = credential_client region = self._config.region sdk_config.endpoint = f"dms-enterprise.{region}.aliyuncs.com" sdk_config.user_agent = "AlibabaCloud-Agent-Skills" self._dms_client = DmsClient(sdk_config) async def list_instances( self, search_key: Optional[str] = None, db_type: Optional[str] = None, env_type: Optional[str] = None, page_number: int = 1, page_size: int = 50, ) -> PagedResult: """Async search for database instances in DMS. Args: search_key: Optional search key. db_type: Optional database type. env_type: Optional environment type. page_number: Page number (default: 1). page_size: Page size (default: 50, max: 100). Returns: PagedResult with list of DmsInstance objects and pagination info. """ if not self._dms_client: await self._init_dms_client() try: page_size = min(page_size, 100) # Max 100 req = dms_models.ListInstancesRequest() req.page_number = page_number req.page_size = page_size if search_key: req.search_key = search_key if db_type: req.db_type = db_type if env_type: req.env_type = env_type resp = await self._dms_client.list_instances_async(req) instance_data = resp.body.to_map() total_count = instance_data.get("TotalCount", 0) result = [] if "InstanceList" in instance_data: instance_list = instance_data["InstanceList"] if "Instance" in instance_list: instances = instance_list["Instance"] if isinstance(instances, list): for item in instances: # Skip DELETED instances state = item.get("State", "") if state == "DELETED": continue resource_id = item.get("EcsInstanceId") or item.get("InstanceResourceId") instance = DmsInstance( instance_id=str(item.get("InstanceId", "")), instance_alias=item.get("InstanceAlias", ""), instance_type=item.get("InstanceType", ""), host=item.get("Host", ""), port=int(item.get("Port", 0)) if item.get("Port") else 0, state=state, env_type=item.get("EnvType", ""), instance_source=item.get("InstanceSource", ""), instance_resource_id=resource_id, ) result.append(instance) return PagedResult( items=result, page_number=page_number, page_size=page_size, total_count=total_count, ) except Exception as e: raise ApiError(f"Failed to list instances: {e}", code="ListInstancesError") FILE:scripts/data_agent/message.py """Message handling for Data Agent. Author: Tinker Created: 2026-03-01 """ from __future__ import annotations import time from typing import Optional, Iterator, AsyncIterator, List from data_agent.client import DataAgentClient, AsyncDataAgentClient from data_agent.sse_client import SSEClient, AsyncSSEClient, SSEEvent from data_agent.models import SessionInfo, ContentBlock, ContentType, AnalysisResult, DataSource from data_agent.exceptions import MessageSendError, ContentFetchError class MessageHandler: """Handles message sending and response retrieval for Data Agent. Uses SSE streaming to receive responses from Data Agent. """ def __init__(self, client: DataAgentClient): """Initialize message handler. Args: client: DataAgentClient instance for API calls. """ self._client = client self._sse_client = SSEClient(client.config) def send_query( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> str: """Send a query and wait for complete response. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Returns: Complete response text from Data Agent. Raises: MessageSendError: If message sending fails. ContentFetchError: If content retrieval fails or times out. """ timeout = timeout or self._client.config.timeout # Send the message try: self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") # Update session usage session.update_last_used() # Get response via SSE streaming try: response = self._sse_client.get_full_response( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ) return response except Exception as e: raise ContentFetchError(f"Failed to get response: {e}") def send_query_with_result( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> AnalysisResult: """Send a query and get detailed result with content blocks. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Returns: AnalysisResult with response text and content blocks. """ timeout = timeout or self._client.config.timeout start_time = time.time() # Send the message try: self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") session.update_last_used() # Collect response via SSE content_blocks = [] response_parts = [] try: for event in self._sse_client.stream_chat_content( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ): block = self._event_to_content_block(event) if block: content_blocks.append(block) if block.content_type == ContentType.TEXT: response_parts.append(block.content) if event.event_type == "SSE_FINISH": break except Exception as e: raise ContentFetchError(f"Failed to get response: {e}") duration_ms = int((time.time() - start_time) * 1000) response_text = "".join(response_parts) return AnalysisResult( query=query, response=response_text, content_blocks=content_blocks, session_id=session.session_id, duration_ms=duration_ms, ) def stream_content( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> Iterator[ContentBlock]: """Send query and stream response content blocks. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Yields: ContentBlock objects as they become available. """ timeout = timeout or self._client.config.timeout # Send the message try: self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") session.update_last_used() # Stream content blocks for event in self._sse_client.stream_chat_content( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ): block = self._event_to_content_block(event) if block: yield block if event.event_type == "SSE_FINISH": break def stream_events( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> Iterator[SSEEvent]: """Send query and stream raw SSE events. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Yields: SSEEvent objects as they arrive. """ timeout = timeout or self._client.config.timeout # Send the message try: self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") session.update_last_used() # Stream raw events yield from self._sse_client.stream_chat_content( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ) def _event_to_content_block(self, event: SSEEvent) -> Optional[ContentBlock]: """Convert SSE event to ContentBlock. Args: event: SSE event to convert. Returns: ContentBlock or None if not applicable. """ # Only process delta events with content if event.event_type == "delta" and event.category == "llm": if event.content: return ContentBlock( content_type=ContentType.TEXT, content=event.content, checkpoint=str(event.checkpoint) if event.checkpoint else None, is_final=False, ) # Process data events (complete content) if event.event_type == "data" and event.category == "think": if event.content: return ContentBlock( content_type=ContentType.TEXT, content=event.content, checkpoint=str(event.checkpoint) if event.checkpoint else None, is_final=True, ) return None class AsyncMessageHandler: """Asynchronous message handler for Data Agent. Uses async SSE streaming for responses. """ def __init__(self, client: AsyncDataAgentClient): """Initialize async message handler. Args: client: AsyncDataAgentClient instance for API calls. """ self._client = client self._sse_client = AsyncSSEClient(client.config) async def send_query( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> str: """Send a query and wait for complete response asynchronously. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Returns: Complete response text from Data Agent. """ timeout = timeout or self._client.config.timeout # Send the message try: await self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") session.update_last_used() # Get response via SSE streaming try: response = await self._sse_client.get_full_response( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ) return response except Exception as e: raise ContentFetchError(f"Failed to get response: {e}") async def send_query_with_result( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> AnalysisResult: """Send a query and get detailed result asynchronously. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Returns: AnalysisResult with response text and content blocks. """ timeout = timeout or self._client.config.timeout start_time = time.time() try: await self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") session.update_last_used() content_blocks = [] response_parts = [] try: async for event in self._sse_client.stream_chat_content( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ): block = self._event_to_content_block(event) if block: content_blocks.append(block) if block.content_type == ContentType.TEXT: response_parts.append(block.content) if event.event_type == "SSE_FINISH": break except Exception as e: raise ContentFetchError(f"Failed to get response: {e}") duration_ms = int((time.time() - start_time) * 1000) response_text = "".join(response_parts) return AnalysisResult( query=query, response=response_text, content_blocks=content_blocks, session_id=session.session_id, duration_ms=duration_ms, ) async def stream_content( self, session: SessionInfo, query: str, timeout: Optional[int] = None, data_source: Optional[DataSource] = None, ) -> AsyncIterator[ContentBlock]: """Send query and stream response content blocks asynchronously. Args: session: Active session to use. query: Natural language query. timeout: Optional timeout in seconds. data_source: Optional DataSource with database metadata. Yields: ContentBlock objects as they become available. """ timeout = timeout or self._client.config.timeout try: await self._client.send_message( agent_id=session.agent_id, session_id=session.session_id, message=query, data_source=data_source, ) except Exception as e: raise MessageSendError(f"Failed to send message: {e}") session.update_last_used() async for event in self._sse_client.stream_chat_content( agent_id=session.agent_id, session_id=session.session_id, timeout=timeout, ): block = self._event_to_content_block(event) if block: yield block if event.event_type == "SSE_FINISH": break def _event_to_content_block(self, event: SSEEvent) -> Optional[ContentBlock]: """Convert SSE event to ContentBlock.""" if event.event_type == "delta" and event.category == "llm": if event.content: return ContentBlock( content_type=ContentType.TEXT, content=event.content, checkpoint=str(event.checkpoint) if event.checkpoint else None, is_final=False, ) if event.event_type == "data" and event.category == "think": if event.content: return ContentBlock( content_type=ContentType.TEXT, content=event.content, checkpoint=str(event.checkpoint) if event.checkpoint else None, is_final=True, ) return None FILE:scripts/data_agent/models.py """Data models for Data Agent SDK. Author: Tinker Created: 2026-03-01 """ from dataclasses import dataclass, field from datetime import datetime from typing import Optional from enum import Enum from .api_adapter import APIAdapter class SessionStatus(str, Enum): """Session status enumeration.""" CREATING = "CREATING" RUNNING = "RUNNING" IDLE = "IDLE" WAIT_INPUT = "WAIT_INPUT" STOPPED = "STOPPED" FAILED = "FAILED" class ContentType(str, Enum): """Content block type enumeration.""" TEXT = "text" IMAGE = "image" TABLE = "table" CHART = "chart" CODE = "code" @dataclass class SessionInfo: """Information about a Data Agent session.""" agent_id: str session_id: str status: SessionStatus = SessionStatus.CREATING database_id: Optional[str] = None created_at: datetime = field(default_factory=datetime.now) last_used_at: datetime = field(default_factory=datetime.now) request_id: Optional[str] = None def is_running(self) -> bool: """Check if session is in an active/usable state (RUNNING, IDLE or WAIT_INPUT).""" return self.status in (SessionStatus.RUNNING, SessionStatus.IDLE, SessionStatus.WAIT_INPUT) def update_last_used(self) -> None: """Update the last used timestamp.""" self.last_used_at = datetime.now() @dataclass class ContentBlock: """A block of content from Data Agent response.""" content_type: ContentType content: str checkpoint: Optional[str] = None sequence_id: Optional[int] = None is_final: bool = False metadata: dict = field(default_factory=dict) @dataclass class FileInfo: """Information about an uploaded or generated file.""" file_id: str filename: str file_type: str size: int upload_url: Optional[str] = None download_url: Optional[str] = None created_at: datetime = field(default_factory=datetime.now) metadata: dict = field(default_factory=dict) @dataclass class DatabaseSource: """Database source configuration for analysis. This is the legacy model. Use DataSource for SendChatMessage API. """ database_id: str database_type: str # PolarDB, RDS, AnalyticDB, MaxCompute, etc. instance_id: Optional[str] = None schema_name: Optional[str] = None description: Optional[str] = None @dataclass class DataSource: """Data source configuration for SendChatMessage API. Contains complete database metadata required by Data Agent. Attributes: dms_instance_id: DMS instance ID (e.g., 1234567) dms_database_id: DMS database ID (e.g., 12345678) instance_name: RDS instance name (e.g., "rm-xxxxxx") db_name: Database name (e.g., "chinook") tables: List of table names to include table_ids: List of table IDs from ListDataCenterTable engine: Database engine type (e.g., "mysql") region_id: Region ID (e.g., "cn-hangzhou") data_source_type: Data source type, default "database" file_id: File ID for file-based analysis (when data_source_type="FILE") """ dms_instance_id: Optional[int] = None dms_database_id: Optional[int] = None instance_name: str = "" db_name: str = "" tables: list[str] = field(default_factory=list) table_ids: list[str] = field(default_factory=list) engine: str = "mysql" region_id: str = "cn-hangzhou" data_source_type: str = "database" file_id: str = "" def to_api_dict(self) -> dict: """Convert to API request format.""" result = { "DataSourceType": self.data_source_type, } if self.data_source_type == "FILE": # File-based analysis uses remote_data_center type with FileId # Extract filename from file_id path for Database field file_name = self.file_id.split("/")[-1] if "/" in self.file_id else self.file_id result.update({ "DataSourceType": "remote_data_center", "FileId": self.file_id, "Database": file_name, "Tables": [file_name.replace(".csv", "").replace(".xlsx", "").replace(".xls", "")], "RegionId": self.region_id, }) else: # Database analysis needs full connection info result.update({ "DmsInstanceId": str(self.dms_instance_id) if self.dms_instance_id else "", "DmsDatabaseId": str(self.dms_database_id) if self.dms_database_id else "", "FileId": self.instance_name, "DbName": self.db_name, "Database": self.db_name, "Tables": self.tables, "TableIds": self.table_ids, "Engine": self.engine, "RegionId": self.region_id, }) # Apply API parameter formatting with PascalCase return APIAdapter.prepare_request_params(result) @dataclass class ChatMessage: """A chat message in the conversation.""" role: str # "user" or "assistant" content: str timestamp: datetime = field(default_factory=datetime.now) message_id: Optional[str] = None @dataclass class AnalysisResult: """Result of a data analysis query.""" query: str response: str content_blocks: list[ContentBlock] = field(default_factory=list) session_id: Optional[str] = None duration_ms: Optional[int] = None generated_files: list[FileInfo] = field(default_factory=list) FILE:scripts/data_agent/session.py """Session management for Data Agent. Author: Tinker Created: 2026-03-01 """ from __future__ import annotations import asyncio import time from typing import Optional, Dict from datetime import datetime, timedelta from data_agent.client import DataAgentClient, AsyncDataAgentClient from data_agent.models import SessionInfo, SessionStatus from data_agent.exceptions import SessionTimeoutError, SessionNotFoundError class SessionManager: """Manages Data Agent session lifecycle. Handles session creation, status polling, and session reuse for multi-turn conversations. """ def __init__(self, client: DataAgentClient): """Initialize session manager. Args: client: DataAgentClient instance for API calls. """ self._client = client self._active_sessions: Dict[str, SessionInfo] = {} def create_or_reuse( self, session_id: Optional[str] = None, agent_id: Optional[str] = None, database_id: Optional[str] = None, wait_for_running: bool = True, mode: Optional[str] = "ASK_DATA", enable_search: bool = False, file_id: Optional[str] = None, ) -> SessionInfo: """Create a new session or reuse an existing one. Args: session_id: Optional existing session ID to reuse. agent_id: Optional agent ID (used when reusing a session to avoid an extra describe call with an empty agent_id). database_id: Optional database ID to bind to new session. wait_for_running: Whether to wait for session to be RUNNING. mode: Session mode to use when creating a new session. One of "ASK_DATA" (default), "ANALYSIS", "INSIGHT". enable_search: Whether to enable search capability in the session. file_id: Optional file ID for file-based analysis session. Returns: SessionInfo for the active session. Raises: SessionTimeoutError: If session doesn't reach RUNNING state. SessionNotFoundError: If specified session_id is invalid. """ # Try to reuse existing session if session_id: if session_id in self._active_sessions: session = self._active_sessions[session_id] if self.is_session_active(session_id): session.update_last_used() return session # Validate session with API try: session = self._client.describe_session( session_id=session_id, agent_id=agent_id or "", ) # If the session is still in CREATING state, wait for it to become ready if session.status == SessionStatus.CREATING and wait_for_running: session = self.wait_until_running( session_id=session.session_id, agent_id=session.agent_id, max_wait=600, # Increase max_wait to 10 minutes (600 seconds) since backend can be slow ) if session.is_running() or not wait_for_running: self._active_sessions[session_id] = session return session else: raise SessionNotFoundError( f"Session {session_id} exists but is not in RUNNING state (status: {session.status.value})" ) except SessionNotFoundError: raise except Exception as e: raise SessionNotFoundError( f"Session {session_id} not found or inaccessible: {e}" ) from e # Create new session session = self._client.create_session(database_id=database_id, mode=mode, enable_search=enable_search, file_id=file_id) # Only wait if session is not already running if wait_for_running and not session.is_running(): session = self.wait_until_running( session_id=session.session_id, agent_id=session.agent_id, max_wait=600, # Increase max_wait to 10 minutes (600 seconds) since backend can be slow ) self._active_sessions[session.session_id] = session return session def wait_until_running( self, session_id: str, agent_id: str = "", max_wait: int = 120, ) -> SessionInfo: """Wait for session to reach RUNNING state. Args: session_id: The session ID. agent_id: The agent ID (optional). max_wait: Maximum wait time in seconds. Returns: SessionInfo with RUNNING status. Raises: SessionTimeoutError: If session doesn't reach RUNNING state. """ config = self._client.config poll_interval = config.poll_interval start_time = time.time() while True: elapsed = time.time() - start_time if elapsed >= max_wait: raise SessionTimeoutError( f"Session {session_id} did not reach RUNNING state within {max_wait} seconds", session_id=session_id, waited_seconds=int(elapsed), ) session = self._client.describe_session(session_id=session_id, agent_id=agent_id) if session.is_running(): return session elif session.status == SessionStatus.FAILED: raise SessionTimeoutError( f"Session {session_id} failed to start", session_id=session_id, waited_seconds=int(elapsed), ) elif session.status == SessionStatus.STOPPED: raise SessionTimeoutError( f"Session {session_id} was stopped unexpectedly", session_id=session_id, waited_seconds=int(elapsed), ) time.sleep(poll_interval) def is_session_active(self, session_id: str) -> bool: """Check if a session is still active and usable. Args: session_id: The session ID to check. Returns: True if session is active and RUNNING. """ if session_id not in self._active_sessions: return False session = self._active_sessions[session_id] # Consider session stale if not used for 30 minutes if datetime.now() - session.last_used_at > timedelta(minutes=30): return False try: current = self._client.describe_session(session_id=session_id, agent_id=session.agent_id) return current.is_running() except Exception: return False def get_session(self, session_id: str) -> Optional[SessionInfo]: """Get session info from cache. Args: session_id: The session ID. Returns: SessionInfo if found, None otherwise. """ return self._active_sessions.get(session_id) def refresh_session(self, session_id: str) -> SessionInfo: """Refresh session status from API. Args: session_id: The session ID. Returns: Updated SessionInfo. Raises: SessionNotFoundError: If session is not in cache. """ if session_id not in self._active_sessions: raise SessionNotFoundError(f"Session {session_id} not found in cache") session = self._active_sessions[session_id] updated = self._client.describe_session(session_id=session_id, agent_id=session.agent_id) self._active_sessions[session_id] = updated return updated def remove_session(self, session_id: str) -> None: """Remove session from cache. Args: session_id: The session ID to remove. """ self._active_sessions.pop(session_id, None) def list_sessions(self) -> list[SessionInfo]: """List all cached sessions. Returns: List of SessionInfo objects. """ return list(self._active_sessions.values()) def clear_stale_sessions(self, max_age_minutes: int = 30) -> int: """Remove stale sessions from cache. Args: max_age_minutes: Maximum session age in minutes. Returns: Number of sessions removed. """ cutoff = datetime.now() - timedelta(minutes=max_age_minutes) stale_ids = [ sid for sid, session in self._active_sessions.items() if session.last_used_at < cutoff ] for sid in stale_ids: del self._active_sessions[sid] return len(stale_ids) class AsyncSessionManager: """Asynchronous session manager for Data Agent. Provides async/await support for session management operations. """ def __init__(self, client: AsyncDataAgentClient): """Initialize async session manager. Args: client: AsyncDataAgentClient instance for API calls. """ self._client = client self._active_sessions: Dict[str, SessionInfo] = {} async def create_or_reuse( self, session_id: Optional[str] = None, database_id: Optional[str] = None, wait_for_running: bool = True, ) -> SessionInfo: """Create a new session or reuse an existing one asynchronously. Args: session_id: Optional existing session ID to reuse. database_id: Optional database ID to bind to new session. wait_for_running: Whether to wait for session to be RUNNING. Returns: SessionInfo for the active session. """ # Try to reuse existing session if session_id: if session_id in self._active_sessions: session = self._active_sessions[session_id] if await self.is_session_active(session_id): session.update_last_used() return session try: session = await self._client.describe_session( session_id=session_id, agent_id="", ) if session.is_running(): self._active_sessions[session_id] = session return session except Exception: pass # Create new session session = await self._client.create_session(database_id=database_id) # Only wait if session is not already running if wait_for_running and not session.is_running(): session = await self.wait_until_running( session_id=session.session_id, agent_id=session.agent_id, ) self._active_sessions[session.session_id] = session return session async def wait_until_running( self, session_id: str, agent_id: str = "", max_wait: int = 120, ) -> SessionInfo: """Wait for session to reach RUNNING state asynchronously. Args: session_id: The session ID. agent_id: The agent ID (optional). max_wait: Maximum wait time in seconds. Returns: SessionInfo with RUNNING status. Raises: SessionTimeoutError: If session doesn't reach RUNNING state. """ config = self._client.config poll_interval = config.poll_interval start_time = time.time() while True: elapsed = time.time() - start_time if elapsed >= max_wait: raise SessionTimeoutError( f"Session {session_id} did not reach RUNNING state within {max_wait} seconds", session_id=session_id, waited_seconds=int(elapsed), ) session = await self._client.describe_session(session_id=session_id, agent_id=agent_id) if session.is_running(): return session elif session.status in (SessionStatus.FAILED, SessionStatus.STOPPED): raise SessionTimeoutError( f"Session {session_id} failed or stopped", session_id=session_id, waited_seconds=int(elapsed), ) await asyncio.sleep(poll_interval) async def is_session_active(self, session_id: str) -> bool: """Check if a session is still active asynchronously. Args: session_id: The session ID to check. Returns: True if session is active and RUNNING. """ if session_id not in self._active_sessions: return False session = self._active_sessions[session_id] if datetime.now() - session.last_used_at > timedelta(minutes=30): return False try: current = await self._client.describe_session(session_id=session_id, agent_id=session.agent_id) return current.is_running() except Exception: return False def get_session(self, session_id: str) -> Optional[SessionInfo]: """Get session info from cache. Args: session_id: The session ID. Returns: SessionInfo if found, None otherwise. """ return self._active_sessions.get(session_id) async def refresh_session(self, session_id: str) -> SessionInfo: """Refresh session status from API asynchronously. Args: session_id: The session ID. Returns: Updated SessionInfo. Raises: SessionNotFoundError: If session is not in cache. """ if session_id not in self._active_sessions: raise SessionNotFoundError(f"Session {session_id} not found in cache") session = self._active_sessions[session_id] updated = await self._client.describe_session(session_id=session_id, agent_id=session.agent_id) self._active_sessions[session_id] = updated return updated def remove_session(self, session_id: str) -> None: """Remove session from cache. Args: session_id: The session ID to remove. """ self._active_sessions.pop(session_id, None) def list_sessions(self) -> list[SessionInfo]: """List all cached sessions. Returns: List of SessionInfo objects. """ return list(self._active_sessions.values()) FILE:scripts/data_agent/sse_client.py """SSE (Server-Sent Events) streaming client for Data Agent. Author: Tinker Created: 2026-03-01 """ from __future__ import annotations import json import hashlib import hmac import base64 import ssl import time import uuid from datetime import datetime, timezone from urllib.parse import urlencode, quote from typing import Iterator, Optional, Dict, Any, AsyncIterator from dataclasses import dataclass import requests import aiohttp from data_agent.config import DataAgentConfig @dataclass class SSEEvent: """Represents a single SSE event.""" event_type: str data: Dict[str, Any] checkpoint: Optional[int] = None category: Optional[str] = None content: Optional[str] = None content_type: Optional[str] = None # "json" or "str" class AliyunSignerV3: """Signs requests for Alibaba Cloud OpenAPI using Signature Version 3.""" # SHA256 hash of empty string EMPTY_HASH = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" @staticmethod def _sha256_hex(data: str) -> str: """Calculate SHA256 hash and return hex string.""" return hashlib.sha256(data.encode("utf-8")).hexdigest() @staticmethod def _hmac_sha256(key: bytes, data: str) -> bytes: """Calculate HMAC-SHA256.""" return hmac.new(key, data.encode("utf-8"), hashlib.sha256).digest() @staticmethod def sign( access_key_id: str, access_key_secret: str, method: str, host: str, action: str, params: Dict[str, str], body: str = "", security_token: Optional[str] = None, ) -> Dict[str, str]: """Generate signed headers for Alibaba Cloud API V3. Args: access_key_id: Access Key ID. access_key_secret: Access Key Secret. method: HTTP method (GET/POST). host: API host. action: API action name. params: Request parameters. body: Request body (empty for GET requests). security_token: STS Security Token (optional). Returns: Headers dict with authorization. """ # Timestamp in ISO 8601 format timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") nonce = str(uuid.uuid4()) # Hash the request body (empty string hash for no body) hashed_payload = AliyunSignerV3._sha256_hex(body) if body else AliyunSignerV3.EMPTY_HASH # Build canonical request http_method = method.upper() canonical_uri = "/" # Query string (sorted) query_params = {"Action": action, "Version": "2025-04-14", **params} sorted_query = sorted(query_params.items()) canonical_query_string = "&".join( f"{quote(k, safe='')}={quote(str(v), safe='')}" for k, v in sorted_query ) # Headers to sign headers_to_sign = { "host": host, "x-acs-action": action, "x-acs-content-sha256": hashed_payload, "x-acs-date": timestamp, "x-acs-signature-nonce": nonce, "x-acs-version": "2025-04-14", } # Add security token to headers if provided if security_token: headers_to_sign["x-acs-security-token"] = security_token # Canonical headers (sorted, lowercase) sorted_headers = sorted(headers_to_sign.items()) canonical_headers = "\n".join(f"{k}:{v}" for k, v in sorted_headers) + "\n" signed_headers = ";".join(k for k, _ in sorted_headers) # Build canonical request string canonical_request = ( f"{http_method}\n" f"{canonical_uri}\n" f"{canonical_query_string}\n" f"{canonical_headers}\n" f"{signed_headers}\n" f"{hashed_payload}" ) # Hash the canonical request hashed_canonical_request = AliyunSignerV3._sha256_hex(canonical_request) # String to sign algorithm = "ACS3-HMAC-SHA256" string_to_sign = f"{algorithm}\n{hashed_canonical_request}" # Calculate signature signature = hmac.new( access_key_secret.encode("utf-8"), string_to_sign.encode("utf-8"), hashlib.sha256, ).hexdigest() # Build authorization header authorization = ( f"{algorithm} Credential={access_key_id}," f"SignedHeaders={signed_headers}," f"Signature={signature}" ) result = { "Authorization": authorization, "Host": host, "x-acs-action": action, "x-acs-content-sha256": hashed_payload, "x-acs-date": timestamp, "x-acs-signature-nonce": nonce, "x-acs-version": "2025-04-14", "Accept": "text/event-stream", } # Add security token to result headers if provided if security_token: result["x-acs-security-token"] = security_token return result # Keep old signer for backward compatibility AliyunSigner = AliyunSignerV3 class SSEClient: """Client for handling SSE streaming responses from Data Agent.""" def __init__(self, config: DataAgentConfig): """Initialize SSE client. Args: config: Data Agent configuration. """ self._config = config self._base_url = f"https://{config.endpoint}" def _do_stream( self, agent_id: str, session_id: str, timeout: int, checkpoint: Optional[int] = None, ) -> Iterator[SSEEvent]: """Single-connection streaming attempt.""" # Check authentication type if hasattr(self._config, 'api_key') and self._config.api_key: # API_KEY authentication return self._do_stream_with_api_key(agent_id, session_id, timeout, checkpoint) else: # AK/SK authentication via default credential chain return self._do_stream_with_ak_sk(agent_id, session_id, timeout, checkpoint) def _do_stream_with_api_key( self, agent_id: str, session_id: str, timeout: int, checkpoint: Optional[int] = None, ) -> Iterator[SSEEvent]: """Stream with API_KEY authentication.""" # Determine endpoint based on action type (GetChatContent is a data plane API) base_endpoint = f"dataagent-stream-{self._config.region}.aliyuncs.com/apikey" url = f"https://{base_endpoint}" # Build request body body = { 'Action': 'GetChatContent', 'Version': '2025-04-14', 'RegionId': self._config.region, 'AgentId': agent_id, 'SessionId': session_id, } if checkpoint is not None: body['Checkpoint'] = str(checkpoint) # Set up headers with API_KEY headers = { 'x-api-key': self._config.api_key, 'Content-Type': 'application/json', 'Accept': 'application/json', 'User-Agent': 'AlibabaCloud-Agent-Skills', } with requests.post( url, stream=True, timeout=timeout, headers=headers, json=body, ) as response: if not response.ok: request_id = response.headers.get("x-acs-request-id", "") body_preview = "" try: body_preview = response.text[:500] except Exception: pass raise requests.exceptions.HTTPError( f"{response.status_code} {response.reason} for {response.url}" f"\n Request-Id: {request_id}" + (f"\n Body: {body_preview}" if body_preview else ""), response=response, ) buffer = "" event_type = "" for chunk in response.iter_content(chunk_size=1024, decode_unicode=True): if not chunk: continue buffer += chunk # Process complete lines from buffer while "\n" in buffer: line, buffer = buffer.split("\n", 1) line = line.strip() if not line: continue if line.startswith("event:"): event_part = line[6:].strip() if " at " in event_part: event_type = event_part.split(" at ")[0] else: event_type = event_part.split()[0] if event_part else "" elif line.startswith("data:"): data_str = line[5:].strip() if data_str and event_type: event = self._parse_event(event_type, data_str) yield event if event.event_type == "SSE_FINISH": return def _do_stream_with_ak_sk( self, agent_id: str, session_id: str, timeout: int, checkpoint: Optional[int] = None, ) -> Iterator[SSEEvent]: """Stream with AK/SK authentication via default credential chain.""" from alibabacloud_credentials.client import Client as CredentialClient credential_client = CredentialClient() credential = credential_client.get_credential() access_key_id = credential.access_key_id access_key_secret = credential.access_key_secret security_token = credential.security_token params = { "AgentId": agent_id, "SessionId": session_id, } if checkpoint is not None: params["Checkpoint"] = str(checkpoint) host = self._config.endpoint headers = AliyunSignerV3.sign( access_key_id, access_key_secret, "POST", host, "GetChatContent", params, security_token=security_token if security_token else None, ) headers["User-Agent"] = "AlibabaCloud-Agent-Skills" # Build query string for URL query_params = {"Action": "GetChatContent", "Version": "2025-04-14", **params} query_string = "&".join( f"{quote(k, safe='')}={quote(str(v), safe='')}" for k, v in sorted(query_params.items()) ) url = f"{self._base_url}/?{query_string}" with requests.post( url, stream=True, timeout=timeout, headers=headers, ) as response: if not response.ok: request_id = response.headers.get("x-acs-request-id", "") body_preview = "" try: body_preview = response.text[:500] except Exception: pass raise requests.exceptions.HTTPError( f"{response.status_code} {response.reason} for {response.url}" f"\n Request-Id: {request_id}" + (f"\n Body: {body_preview}" if body_preview else ""), response=response, ) buffer = "" event_type = "" for chunk in response.iter_content(chunk_size=1024, decode_unicode=True): if not chunk: continue buffer += chunk # Process complete lines from buffer while "\n" in buffer: line, buffer = buffer.split("\n", 1) line = line.strip() if not line: continue if line.startswith("event:"): # Extract event type (before " at " timestamp) event_part = line[6:].strip() if " at " in event_part: event_type = event_part.split(" at ")[0] else: event_type = event_part.split()[0] if event_part else "" elif line.startswith("data:"): data_str = line[5:].strip() if data_str and event_type: event = self._parse_event(event_type, data_str) yield event if event.event_type == "SSE_FINISH": return def stream_chat_content( self, agent_id: str, session_id: str, timeout: int = 120, checkpoint: Optional[int] = None, max_retries: int = 3, ) -> Iterator[SSEEvent]: """Stream chat content from Data Agent using SSE. Automatically reconnects from the last checkpoint on SSL/connection errors, enabling resilience during long-running ANALYSIS sessions. Args: agent_id: The agent ID. session_id: The session ID. timeout: Request timeout in seconds. checkpoint: Optional starting checkpoint (0 to replay from beginning). max_retries: Maximum number of reconnection attempts on error. Yields: SSEEvent objects as they arrive. """ # Retryable network/SSL exception types _retryable = ( ssl.SSLError, requests.exceptions.SSLError, requests.exceptions.ChunkedEncodingError, requests.exceptions.ConnectionError, ) current_checkpoint = checkpoint retry_count = 0 while True: try: for event in self._do_stream(agent_id, session_id, timeout, current_checkpoint): # Track latest checkpoint for potential reconnection if event.checkpoint is not None: current_checkpoint = event.checkpoint yield event if event.event_type == "SSE_FINISH": return return # Clean finish without SSE_FINISH except _retryable as exc: retry_count += 1 if retry_count > max_retries: raise wait = min(2 ** retry_count, 30) # 2s, 4s, 8s … capped at 30s time.sleep(wait) # Loop continues with updated current_checkpoint def _parse_event(self, event_type: str, data_str: str) -> SSEEvent: """Parse an SSE event. Args: event_type: The event type. data_str: The data JSON string. Returns: Parsed SSEEvent. """ try: data = json.loads(data_str) except json.JSONDecodeError: data = {"content": data_str} return SSEEvent( event_type=event_type, data=data, checkpoint=data.get("checkpoint"), category=data.get("category"), content=data.get("content", ""), content_type=data.get("content_type"), ) # Categories to skip when collecting data events _SKIP_DATA_CATEGORIES = {"tool_call_choices", "metric_agent_config"} def get_full_response( self, agent_id: str, session_id: str, timeout: int = 120, ) -> str: """Get complete response by collecting all SSE events. Collects all `data` events (excluding tool_call_choices and metric_agent_config), which carry the primary analysis results. Also collects `delta/llm` events as supplementary streaming text. Args: agent_id: The agent ID. session_id: The session ID. timeout: Request timeout in seconds. Returns: Complete response text. """ content_parts = [] for event in self.stream_chat_content(agent_id, session_id, timeout): # Collect all delta events (llm / tool_call_response / output_conclusion etc.) if event.event_type == "delta": if event.content: content_parts.append(event.content) # Collect data events except filtered categories elif event.event_type == "data": if event.category not in self._SKIP_DATA_CATEGORIES: if event.content: content_parts.append(event.content) if event.event_type == "SSE_FINISH": break return "".join(content_parts) class AsyncSSEClient: """Async client for handling SSE streaming responses.""" def __init__(self, config: DataAgentConfig): """Initialize async SSE client. Args: config: Data Agent configuration. """ self._config = config self._base_url = f"https://{config.endpoint}" async def _async_do_stream( self, agent_id: str, session_id: str, timeout: int, checkpoint: Optional[int] = None, ) -> AsyncIterator[SSEEvent]: """Single-connection async streaming attempt.""" # Check authentication type if hasattr(self._config, 'api_key') and self._config.api_key: # API_KEY authentication async for event in self._async_do_stream_with_api_key(agent_id, session_id, timeout, checkpoint): yield event else: # AK/SK authentication via default credential chain async for event in self._async_do_stream_with_ak_sk(agent_id, session_id, timeout, checkpoint): yield event async def _async_do_stream_with_api_key( self, agent_id: str, session_id: str, timeout: int, checkpoint: Optional[int] = None, ) -> AsyncIterator[SSEEvent]: """Async stream with API_KEY authentication.""" # Determine endpoint based on action type (GetChatContent is a data plane API) base_endpoint = f"dataagent-stream-{self._config.region}.aliyuncs.com/apikey" url = f"https://{base_endpoint}" # Build request body body = { 'Action': 'GetChatContent', 'Version': '2025-04-14', 'RegionId': self._config.region, 'AgentId': agent_id, 'SessionId': session_id, } if checkpoint is not None: body['Checkpoint'] = str(checkpoint) # Set up headers with API_KEY headers = { 'x-api-key': self._config.api_key, 'Content-Type': 'application/json', 'Accept': 'application/json', 'User-Agent': 'AlibabaCloud-Agent-Skills', } async with aiohttp.ClientSession() as session: async with session.post( url, timeout=aiohttp.ClientTimeout(total=timeout), headers=headers, json=body, ) as response: response.raise_for_status() buffer = "" event_type = "" async for chunk in response.content.iter_chunked(1024): if not chunk: continue buffer += chunk.decode("utf-8") # Process complete lines from buffer while "\n" in buffer: line, buffer = buffer.split("\n", 1) line = line.strip() if not line: continue if line.startswith("event:"): event_part = line[6:].strip() if " at " in event_part: event_type = event_part.split(" at ")[0] else: event_type = event_part.split()[0] if event_part else "" elif line.startswith("data:"): data_str = line[5:].strip() if data_str and event_type: event = self._parse_event(event_type, data_str) yield event if event.event_type == "SSE_FINISH": return async def _async_do_stream_with_ak_sk( self, agent_id: str, session_id: str, timeout: int, checkpoint: Optional[int] = None, ) -> AsyncIterator[SSEEvent]: """Async stream with AK/SK authentication via default credential chain.""" from alibabacloud_credentials.client import Client as CredentialClient credential_client = CredentialClient() credential = credential_client.get_credential() access_key_id = credential.access_key_id access_key_secret = credential.access_key_secret security_token = credential.security_token params = { "AgentId": agent_id, "SessionId": session_id, } if checkpoint is not None: params["Checkpoint"] = str(checkpoint) host = self._config.endpoint headers = AliyunSignerV3.sign( access_key_id, access_key_secret, "POST", host, "GetChatContent", params, security_token=security_token if security_token else None, ) headers["User-Agent"] = "AlibabaCloud-Agent-Skills" # Build query string for URL query_params = {"Action": "GetChatContent", "Version": "2025-04-14", **params} query_string = "&".join( f"{quote(k, safe='')}={quote(str(v), safe='')}" for k, v in sorted(query_params.items()) ) url = f"{self._base_url}/?{query_string}" async with aiohttp.ClientSession() as session: async with session.post( url, timeout=aiohttp.ClientTimeout(total=timeout), headers=headers, ) as response: response.raise_for_status() buffer = "" event_type = "" async for chunk in response.content.iter_chunked(1024): if not chunk: continue buffer += chunk.decode("utf-8") # Process complete lines from buffer while "\n" in buffer: line, buffer = buffer.split("\n", 1) line = line.strip() if not line: continue if line.startswith("event:"): event_part = line[6:].strip() if " at " in event_part: event_type = event_part.split(" at ")[0] else: event_type = event_part.split()[0] if event_part else "" elif line.startswith("data:"): data_str = line[5:].strip() if data_str and event_type: event = self._parse_event(event_type, data_str) yield event if event.event_type == "SSE_FINISH": return async def stream_chat_content( self, agent_id: str, session_id: str, timeout: int = 120, checkpoint: Optional[int] = None, max_retries: int = 3, ) -> AsyncIterator[SSEEvent]: """Stream chat content asynchronously. Automatically reconnects from the last checkpoint on SSL/connection errors, enabling resilience during long-running ANALYSIS sessions. Args: agent_id: The agent ID. session_id: The session ID. timeout: Request timeout in seconds. checkpoint: Optional starting checkpoint (0 to replay from beginning). max_retries: Maximum number of reconnection attempts on error. Yields: SSEEvent objects as they arrive. """ import asyncio _retryable = ( ssl.SSLError, aiohttp.ClientSSLError, aiohttp.ServerDisconnectedError, aiohttp.ClientConnectionError, ) current_checkpoint = checkpoint retry_count = 0 while True: try: async for event in self._async_do_stream( agent_id, session_id, timeout, current_checkpoint ): if event.checkpoint is not None: current_checkpoint = event.checkpoint yield event if event.event_type == "SSE_FINISH": return return # Clean finish without SSE_FINISH except _retryable: retry_count += 1 if retry_count > max_retries: raise wait = min(2 ** retry_count, 30) await asyncio.sleep(wait) # Loop continues with updated current_checkpoint def _parse_event(self, event_type: str, data_str: str) -> SSEEvent: """Parse an SSE event.""" try: data = json.loads(data_str) except json.JSONDecodeError: data = {"content": data_str} return SSEEvent( event_type=event_type, data=data, checkpoint=data.get("checkpoint"), category=data.get("category"), content=data.get("content", ""), content_type=data.get("content_type"), ) # Categories to skip when collecting data events _SKIP_DATA_CATEGORIES = {"tool_call_choices", "metric_agent_config"} async def get_full_response( self, agent_id: str, session_id: str, timeout: int = 120, ) -> str: """Get complete response asynchronously. Collects all `data` events (excluding tool_call_choices and metric_agent_config), which carry the primary analysis results. Also collects `delta/llm` events as supplementary streaming text. Args: agent_id: The agent ID. session_id: The session ID. timeout: Request timeout in seconds. Returns: Complete response text. """ content_parts = [] async for event in self.stream_chat_content(agent_id, session_id, timeout): # Collect all delta events (llm / tool_call_response / output_conclusion etc.) if event.event_type == "delta": if event.content: content_parts.append(event.content) # Collect data events except filtered categories elif event.event_type == "data": if event.category not in self._SKIP_DATA_CATEGORIES: if event.content: content_parts.append(event.content) if event.event_type == "SSE_FINISH": break return "".join(content_parts) FILE:scripts/data_agent_cli.py #!/usr/bin/env python3 """Data Agent Unified CLI Tool - Entry Point. See ``cli/`` package for implementation details. Author: Tinker Created: 2026-03-03 """ import sys from pathlib import Path # Ensure local packages (data_agent, cli) are importable _SCRIPT_DIR = Path(__file__).resolve().parent if str(_SCRIPT_DIR) not in sys.path: sys.path.insert(0, str(_SCRIPT_DIR)) # Re-export public API for backward compatibility (tests patch "data_agent_cli.XXX") from cli import * # noqa: F401,F403 from cli.parser import main if __name__ == "__main__": main() FILE:scripts/requirements.txt # Data Agent CLI 依赖 # 推荐使用虚拟环境安装: # python3 -m venv venv && source venv/bin/activate # pip install -r requirements.txt # 阿里云 SDK alibabacloud-tea-openapi>=0.3.0 alibabacloud-dms20250414>=1.0.0 alibabacloud-dms-enterprise20181101>=1.0.0 alibabacloud-openapi-util>=0.2.0 # HTTP 客户端 requests>=2.28.0 aiohttp>=3.9.0 # 环境变量 python-dotenv>=1.0.0

Alibabacloud Data Agent Skill

Comments (0)