#!/usr/bin/env python3 """ etlworks_agent.py — Python client for the Etlworks AI Agent Scripting API. This module provides the EtlworksAgent class for programmatic access to: - Individual agent tools (search KB, run CLI commands, search/import templates) - Stateless chat (one-shot question/answer) - Stateful sessions (multi-turn conversations) - Streaming responses (real-time token delivery) Usage as a library: from etlworks_agent import EtlworksAgent agent = EtlworksAgent("https://app.etlworks.com", api_key="your-key") # List tools tools = agent.list_tools() # Execute a tool result = agent.execute_tool("search_knowledge_base", {"query": "REST connection"}) # One-shot chat response = agent.chat("How do I create a connection?") # Streaming chat for event in agent.chat_stream("How do I schedule a flow?"): if event["type"] == "token": print(event["content"], end="", flush=True) # Multi-turn session session = agent.create_session() r1 = agent.session_chat(session, "How do I create a REST connection?") r2 = agent.session_chat(session, "Now schedule it to run hourly") messages = agent.session_messages(session) Usage as a CLI: export ETLWORKS_URL="https://app.etlworks.com" export ETLWORKS_API_KEY="your-key" python etlworks_agent.py tools python etlworks_agent.py tool search_knowledge_base '{"query": "REST"}' python etlworks_agent.py chat "How do I create a connection?" python etlworks_agent.py chat-stream "How do I schedule a flow?" Requirements: requests (pip install requests) """ import json import os import sys import requests class EtlworksAgent: """Client for the Etlworks AI Agent Scripting API.""" def __init__(self, base_url=None, api_key=None): """ Initialize the client. Args: base_url: Etlworks instance URL (default: ETLWORKS_URL env var) api_key: API key (default: ETLWORKS_API_KEY env var) """ self.base_url = (base_url or os.environ.get("ETLWORKS_URL", "http://localhost:8080")).rstrip("/") self.api_key = api_key or os.environ.get("ETLWORKS_API_KEY", "") if not self.api_key: raise ValueError( "API key is required. Pass api_key= or set ETLWORKS_API_KEY env var.\n" "Generate an API key in Etlworks: Settings > Users > Edit user > API Key" ) self.api_base = self.base_url + "/rest/v1/ai-agent/api" self.session_http = requests.Session() self.session_http.headers.update({ "Authorization": "Bearer " + self.api_key, "Content-Type": "application/json", }) # --------------------------------------------------------------- # Tool operations # --------------------------------------------------------------- def list_tools(self): """List all available agent tools with their parameter schemas.""" resp = self.session_http.get(self.api_base + "/tools") resp.raise_for_status() return resp.json() def execute_tool(self, tool_name, arguments=None): """ Execute a specific tool. Args: tool_name: Tool name (e.g. "search_knowledge_base", "host_cli") arguments: Dict of tool arguments (matches the tool's parameter schema) Returns: Dict with "tool", "result", "success" keys """ resp = self.session_http.post( self.api_base + "/tools/" + tool_name + "/execute", json=arguments or {}, ) resp.raise_for_status() return resp.json() # --------------------------------------------------------------- # Stateless chat # --------------------------------------------------------------- def chat(self, message, include_tools=True, system_prompt=None): """ Send a one-shot message and get a complete response. Args: message: The user message include_tools: Whether to enable tool use (default: True) system_prompt: Optional custom system prompt override Returns: Dict with "response", "session_id", "tools_used", "usage" keys """ body = {"message": message, "include_tools": include_tools} if system_prompt: body["system_prompt"] = system_prompt resp = self.session_http.post(self.api_base + "/chat", json=body) resp.raise_for_status() return resp.json() def chat_stream(self, message, include_tools=True, system_prompt=None): """ Send a one-shot message and get a streaming response. Yields dicts with event data. Event types: - {"type": "session", "session_id": "..."} - {"type": "token", "content": "..."} - {"type": "phase", "phase": "...", "message": "..."} - {"type": "tool_start", "tool": "...", "message": "..."} - {"type": "tool_end", "tool": "..."} - {"type": "usage", "prompt_tokens": N, "completion_tokens": N} - {"type": "end"} - {"type": "error", "message": "...", "retryable": bool} Args: message: The user message include_tools: Whether to enable tool use system_prompt: Optional custom system prompt override Yields: Dict for each SSE event """ body = {"message": message, "include_tools": include_tools} if system_prompt: body["system_prompt"] = system_prompt return self._stream_sse(self.api_base + "/chat/stream", body) # --------------------------------------------------------------- # Stateful sessions # --------------------------------------------------------------- def create_session(self): """ Create a new multi-turn chat session. Returns: Session ID string """ resp = self.session_http.post(self.api_base + "/sessions") resp.raise_for_status() return resp.json()["session_id"] def session_chat(self, session_id, message, include_tools=True): """ Send a message in a session and get a complete response. The agent remembers all prior messages in the session. Args: session_id: Session ID from create_session() message: The user message include_tools: Whether to enable tool use Returns: Dict with "response", "session_id", "tools_used", "usage" keys """ body = {"message": message, "include_tools": include_tools} resp = self.session_http.post( self.api_base + "/sessions/" + session_id + "/chat", json=body, ) resp.raise_for_status() return resp.json() def session_chat_stream(self, session_id, message, include_tools=True): """ Send a message in a session and get a streaming response. Args: session_id: Session ID from create_session() message: The user message include_tools: Whether to enable tool use Yields: Dict for each SSE event """ body = {"message": message, "include_tools": include_tools} return self._stream_sse( self.api_base + "/sessions/" + session_id + "/chat/stream", body, ) def session_messages(self, session_id): """ Get conversation history for a session. Returns: List of {"role": "user"|"assistant", "content": "..."} dicts """ resp = self.session_http.get( self.api_base + "/sessions/" + session_id + "/messages", ) resp.raise_for_status() return resp.json() # --------------------------------------------------------------- # Internal # --------------------------------------------------------------- def _stream_sse(self, url, body): """Stream SSE events from a POST endpoint.""" resp = self.session_http.post(url, json=body, stream=True) resp.raise_for_status() for line in resp.iter_lines(decode_unicode=True): if line and line.startswith("data: "): try: yield json.loads(line[6:]) except json.JSONDecodeError: pass # --------------------------------------------------------------- # CLI interface # --------------------------------------------------------------- def _print_stream(events): """Print streaming events to stdout.""" for event in events: t = event.get("type", "") if t == "token": print(event.get("content", ""), end="", flush=True) elif t == "phase": print("\n[%s] %s" % (event.get("phase", ""), event.get("message", ""))) elif t == "tool_start": print("\n[tool] %s: %s" % (event.get("tool", ""), event.get("message", ""))) elif t == "tool_end": print("[tool] %s: done" % event.get("tool", "")) elif t == "error": print("\n[ERROR] %s" % event.get("message", "")) elif t == "usage": print("[usage] prompt=%s completion=%s" % ( event.get("prompt_tokens", ""), event.get("completion_tokens", ""))) elif t == "end": print("\n--- Done ---") elif t == "session": pass # session ID header, not user-visible def main(): args = sys.argv[1:] if not args or args[0] in ("help", "--help", "-h"): print("Etlworks AI Agent Scripting API -- Python Client") print() print("Environment:") print(" ETLWORKS_URL Base URL (default: http://localhost:8080)") print(" ETLWORKS_API_KEY API key (required)") print() print("Commands:") print(" tools List available tools") print(" tool Execute a tool") print(" chat One-shot chat") print(" chat-stream Streaming chat") print(" session-create Create session (prints ID)") print(" session-chat Chat in session") print(" session-chat-stream Streaming chat in session") print(" session-messages Get session history") sys.exit(0) agent = EtlworksAgent() cmd = args[0] if cmd == "tools": result = agent.list_tools() print(json.dumps(result, indent=2)) elif cmd == "tool": name = args[1] tool_args = json.loads(args[2]) if len(args) > 2 else {} result = agent.execute_tool(name, tool_args) print(json.dumps(result, indent=2)) elif cmd == "chat": message = args[1] result = agent.chat(message) print(json.dumps(result, indent=2)) elif cmd == "chat-stream": message = args[1] _print_stream(agent.chat_stream(message)) elif cmd == "session-create": session_id = agent.create_session() print(session_id) elif cmd == "session-chat": session_id, message = args[1], args[2] result = agent.session_chat(session_id, message) print(json.dumps(result, indent=2)) elif cmd == "session-chat-stream": session_id, message = args[1], args[2] _print_stream(agent.session_chat_stream(session_id, message)) elif cmd == "session-messages": session_id = args[1] messages = agent.session_messages(session_id) print(json.dumps(messages, indent=2)) else: print("Unknown command: %s" % cmd, file=sys.stderr) sys.exit(1) if __name__ == "__main__": main()