From 2b8e637298acee145a6db8f67f88e3eeaeb3be43 Mon Sep 17 00:00:00 2001 From: Fede Luzzi Date: Mon, 1 Jun 2026 17:49:19 -0300 Subject: [PATCH] added AI support for yaml/run --- connpy/_version.py | 2 +- connpy/ai.py | 317 ++++++++++- connpy/cli/ai_handler.py | 4 +- connpy/cli/run_handler.py | 288 +++++++++- connpy/completion.py | 5 + connpy/connapp.py | 3 + connpy/grpc_layer/connpy_pb2.py | 42 +- connpy/grpc_layer/connpy_pb2_grpc.py | 172 ++++-- connpy/grpc_layer/server.py | 107 +++- connpy/grpc_layer/stubs.py | 150 ++++- connpy/printer.py | 2 +- connpy/proto/connpy.proto | 14 +- connpy/services/ai_service.py | 34 ++ connpy/services/execution_service.py | 50 -- connpy/tests/test_ai.py | 9 + connpy/tests/test_cli_run_ai.py | 136 +++++ connpy/tests/test_run_ai.py | 296 ++++++++++ docs/connpy/cli/ai_handler.html | 8 +- docs/connpy/cli/run_handler.html | 593 +++++++++++++++++++- docs/connpy/grpc_layer/connpy_pb2_grpc.html | 396 +++++++++---- docs/connpy/grpc_layer/server.html | 111 +++- docs/connpy/grpc_layer/stubs.html | 405 ++++++------- docs/connpy/index.html | 114 ++-- docs/connpy/services/ai_service.html | 103 +++- docs/connpy/services/execution_service.html | 111 +--- docs/connpy/services/index.html | 214 ++++--- 26 files changed, 2885 insertions(+), 801 deletions(-) create mode 100644 connpy/tests/test_cli_run_ai.py create mode 100644 connpy/tests/test_run_ai.py diff --git a/connpy/_version.py b/connpy/_version.py index 0f607a5..79a961b 100644 --- a/connpy/_version.py +++ b/connpy/_version.py @@ -1 +1 @@ -__version__ = "6.0.0" +__version__ = "6.0.1" diff --git a/connpy/ai.py b/connpy/ai.py index a3e3175..62aeeda 100755 --- a/connpy/ai.py +++ b/connpy/ai.py @@ -114,6 +114,7 @@ class ai: self.confirm_handler = confirm_handler or self._local_confirm_handler self.trusted_session = trust # Trust mode for the entire session self.interrupted = False + self.one_shot = kwargs.get("one_shot", False) # 1. Cargar configuración genérica con herencia/merge global @@ -285,10 +286,13 @@ class ai: @property def architect_system_prompt(self): """Build architect system prompt with plugin extensions.""" + prompt = self._architect_base_prompt + if getattr(self, "one_shot", False): + prompt += "\n\nCRITICAL 1-SHOT DIAGNOSTICS DIRECTIVE:\nYou are running in a 1-shot offline diagnostics mode. There is no active conversation loop, and you are NOT conversing with a Network Engineer. You MUST deliver your complete strategic analysis immediately and directly to the user. Do not suggest or attempt to delegate/return control to the engineer." if self.architect_prompt_extensions: extensions = "\n".join(self.architect_prompt_extensions) - return self._architect_base_prompt + f"\n\nPlugin Capabilities:\n{extensions}" - return self._architect_base_prompt + return prompt + f"\n\nPlugin Capabilities:\n{extensions}" + return prompt def register_ai_tool(self, tool_definition, handler, target="engineer", engineer_prompt=None, architect_prompt=None, status_formatter=None): """Register an external tool for the AI system. @@ -880,6 +884,8 @@ class ai: {"type": "function", "function": {"name": "return_to_engineer", "description": "Return control to the Engineer. Use this when your strategic analysis is complete and the Engineer should handle the rest of the conversation.", "parameters": {"type": "object", "properties": {"summary": {"type": "string", "description": "Brief summary of your analysis to hand over to the Engineer."}}, "required": ["summary"]}}}, {"type": "function", "function": {"name": "manage_memory_tool", "description": "Saves information to long-term memory. MANDATORY: Only use this if the user explicitly asks to remember or save something.", "parameters": {"type": "object", "properties": {"content": {"type": "string"}, "action": {"type": "string", "enum": ["append", "replace"]}}, "required": ["content"]}}} ] + if getattr(self, "one_shot", False): + base_tools = [t for t in base_tools if t["function"]["name"] not in ("delegate_to_engineer", "return_to_engineer")] all_tools = base_tools + self.external_architect_tools seen_names = set() @@ -1624,3 +1630,310 @@ Node: {node_name}""" @MethodHook def confirm(self, user_input): return True + + +PLAYBOOK_BUILDER_SYSTEM_PROMPT = """ +You are a Connpy Playbook Builder Agent, a specialist in creating structured Connpy automation playbooks in YAML format. +Your primary mission is to help the user build, refine, and validate playbooks. + +You MUST follow the Connpy canonical playbook format strictly: +The playbook MUST always use the `tasks[]` array structure as the root key, where each task is sequential and independent. + +Connpy YAML Playbook Canonical Schema: +--- +tasks: +- name: "Task Description" + action: 'run' # Can be 'run' or 'test'. Mandatory. + nodes: # List of nodes filter or regular expressions to work on. Mandatory. Can be a string or array of strings. Supports regex (e.g. 'router.*@office' to match all routers in the 'office' folder). + - 'router1@office' + - 'router.*@office' # Regex filters are fully supported to match multiple nodes dynamically. + - '@aws' + commands: # List of CLI commands to execute. Mandatory. + - 'show version' + variables: # Key-value pairs for variables replacement in commands and expected. Optional. + __global__: # Global variables fallback. Optional. + key: value + node_name@folder: # Node-specific variables. Optional. + key: value + output: stdout # Mandatory. Output configuration. Choices: 'stdout', 'null', or a folder path like '/path/to/folder'. + options: # Execution options. Optional. + prompt: 'regex_prompt' # Optional prompt to expect. + parallel: 10 # Optional number of parallel threads. Default 10. + timeout: 20 # Optional execution timeout in seconds. Default 20. + +- name: "Verification Task" + action: 'test' + nodes: + - 'router1@office' + commands: + - 'ping 10.100.100.1' + expected: '!' # Expected text pattern to search in output. Mandatory ONLY for 'test' action. + +Connpy Variable Templating & Usage: +- Variables defined under the `variables` key (either globally under `__global__` or for specific nodes) are used in commands or expected output by surrounding the variable name with single curly braces: `{variable_name}`. +- Example: If you define a variable `ip` with a value of `10.100.100.1`, you use it in commands as `'ping {ip}'`. +- Recommendation (Important): Variables are not limited to simple words or values. You can define entire CLI commands as variables to abstract vendor-specific syntax! This is highly recommended when executing the same logical operation across different operating systems (OS) or vendors. + - Example: You can define `show_interface_cmd` under a specific node's variables to be `'show ip interface brief'` for Cisco, and `'show interfaces terse'` for Juniper, and then write a single generic command under `commands`: + `- '{show_interface_cmd}'` + +Guidelines: +1. When the user requests a playbook, you should guide them and output the YAML. +2. IMPORTANT: You have access to the `list_nodes` tool. Proactively use it to inspect the user's real inventory. This allows you to discover correct node names, folders, or device tags, and construct precise regex filters for the `nodes` field based on real assets. +3. IMPORTANT: Before presenting the playbook, you MUST call the `validate_playbook` tool with the YAML to let the backend check for syntax and schema correctness. +4. If `validate_playbook` returns errors, fix them in your YAML and validate again before responding to the user. +5. When the playbook is complete, validated, and the user approves it, you MUST call the `return_playbook` tool to return the final YAML. +6. All text responses must be in the same language the user uses in their prompt. +""" + +PLAYBOOK_BUILDER_TOOLS = [ + { + "type": "function", + "function": { + "name": "list_nodes", + "description": "[Universal Platform] Lists available nodes in the inventory. Use this to discover device names, folders, or operating systems to build proper regex filters.", + "parameters": { + "type": "OBJECT", + "properties": { + "filter_pattern": { + "type": "STRING", + "description": "Regex or pattern to filter nodes (e.g. '.*', 'border.*', '@office')." + } + } + } + } + }, + { + "type": "function", + "function": { + "name": "validate_playbook", + "description": "Validates the Connpy YAML playbook structure, syntax, and schema correctness with the backend.", + "parameters": { + "type": "OBJECT", + "properties": { + "playbook_yaml": { + "type": "STRING", + "description": "The YAML content of the playbook to validate." + } + }, + "required": ["playbook_yaml"] + } + } + }, + { + "type": "function", + "function": { + "name": "return_playbook", + "description": "Returns the final validated YAML playbook to the calling application when the user is satisfied.", + "parameters": { + "type": "OBJECT", + "properties": { + "playbook_yaml": { + "type": "STRING", + "description": "The final YAML content of the playbook." + } + }, + "required": ["playbook_yaml"] + } + } + } +] + +class PlaybookBuilderAgent: + """Specialized AI agent for building, validating, and generating Connpy YAML playbooks.""" + + def __init__(self, config, console=None, confirm_handler=None, trust=False, **kwargs): + self.config = config + self.console = console or printer.console + self.interrupted = False + + # Load AI configuration + if hasattr(self.config, "get_effective_setting"): + aiconfig = self.config.get_effective_setting("ai", {}) + else: + aiconfig = self.config.config.get("ai", {}) if hasattr(self.config, "config") else {} + + # Default model for technical tasks + self.model = kwargs.get("engineer_model") or aiconfig.get("engineer_model") or "gemini/gemini-3.1-flash-lite" + self.key = kwargs.get("engineer_api_key") or aiconfig.get("engineer_api_key") + self.auth = kwargs.get("engineer_auth") or aiconfig.get("engineer_auth") or {} + if self.key and "api_key" not in self.auth: + self.auth = self.auth.copy() + self.auth["api_key"] = self.key + + def validate_playbook(self, playbook_yaml: str) -> dict: + """Sintactical and schema validation of Connpy Playbook YAML.""" + import yaml + try: + # 1. Parse YAML + data = yaml.load(playbook_yaml, Loader=yaml.FullLoader) + except Exception as e: + return {"valid": False, "error": f"YAML Syntax Error: {e}"} + + # 2. Check structure + if not isinstance(data, dict): + return {"valid": False, "error": "Playbook must be a YAML dictionary."} + + if "tasks" not in data: + return {"valid": False, "error": "Playbook missing mandatory root 'tasks' key."} + + tasks = data["tasks"] + if not isinstance(tasks, list): + return {"valid": False, "error": "'tasks' must be a list of tasks."} + + # 3. Check individual tasks + for idx, task in enumerate(tasks): + if not isinstance(task, dict): + return {"valid": False, "error": f"Task index {idx} must be a dictionary."} + + name = task.get("name", f"Task {idx}") + + # Mandatory fields + mandatory = ["name", "action", "nodes", "commands", "output"] + missing = [field for field in mandatory if field not in task] + if missing: + return {"valid": False, "error": f"Task '{name}' (index {idx}) is missing mandatory fields: {missing}"} + + # Validate nodes field type (supports string regexes or array of string regexes) + nodes = task["nodes"] + if not isinstance(nodes, (str, list)): + return {"valid": False, "error": f"Task '{name}' (index {idx}) 'nodes' must be a string (regex) or a list of strings (regexes)."} + + if isinstance(nodes, list): + for n_idx, node_item in enumerate(nodes): + if not isinstance(node_item, str): + return {"valid": False, "error": f"Task '{name}' (index {idx}) 'nodes' list contains a non-string value at index {n_idx}: {node_item}"} + + action = task["action"] + if action not in ["run", "test"]: + return {"valid": False, "error": f"Task '{name}' (index {idx}) has invalid action '{action}'. Choices are: 'run', 'test'."} + + if action == "test" and "expected" not in task: + return {"valid": False, "error": f"Task '{name}' (index {idx}) has action 'test' but is missing the mandatory 'expected' key."} + + output = task["output"] + if output not in [None, "stdout"] and not output.startswith("/"): + return {"valid": False, "error": f"Task '{name}' (index {idx}) output '{output}' is invalid. Must be 'stdout', 'null' or an absolute path."} + + return {"valid": True, "message": "Playbook schema and syntax is valid."} + + def ask(self, user_input, chat_history=None, status=None, debug=False, chunk_callback=None): + """Standard conversation step with tool loop for PlaybookBuilderAgent.""" + if chat_history is None: + chat_history = [] + + # System prompt and tool definition + system_prompt = PLAYBOOK_BUILDER_SYSTEM_PROMPT + tools = PLAYBOOK_BUILDER_TOOLS + messages = [{"role": "system", "content": system_prompt}] + + for msg in chat_history: + m = msg if isinstance(msg, dict) else msg.copy() + if m.get('role') == 'assistant' and m.get('tool_calls') and m.get('content') == "": + m['content'] = None + messages.append(m) + + messages.append({"role": "user", "content": user_input}) + + final_playbook_yaml = None + iteration = 0 + max_iterations = 10 + + while iteration < max_iterations: + iteration += 1 + + if status: + status.update(f"Playbook Agent is thinking... (step {iteration})") + + # Call LiteLLM completion + from connpy.ai import completion + try: + response = completion( + model=self.model, + messages=messages, + tools=tools, + num_retries=3, + **self.auth + ) + except Exception as e: + return {"response": f"Playbook Agent failed: {str(e)}", "chat_history": messages[1:]} + + resp_msg = response.choices[0].message + msg_dict = resp_msg.model_dump(exclude_none=True) + if msg_dict.get("tool_calls") and msg_dict.get("content") == "": + msg_dict["content"] = None + + messages.append(msg_dict) + + # If the model sends content, stream or yield it + if resp_msg.content: + if chunk_callback: + chunk_callback(resp_msg.content) + elif not resp_msg.tool_calls: + # In direct non-streaming output, print markdown + self.console.print(Markdown(resp_msg.content)) + + if not resp_msg.tool_calls: + break + + for tc in resp_msg.tool_calls: + fn = tc.function.name + args = json.loads(tc.function.arguments) + + if fn == "list_nodes": + filter_pattern = args.get("filter_pattern", ".*") + try: + matched_names = self.config._getallnodes(filter_pattern) + if not matched_names: + obs = "No nodes found matching the filter." + else: + if len(matched_names) <= 5: + matched_data = self.config.getitems(matched_names, extract=True) + res = {} + for name, data in matched_data.items(): + os_tag = "unknown" + if isinstance(data, dict): + ts = data.get("tags") + if isinstance(ts, dict): os_tag = ts.get("os", "unknown") + res[name] = {"os": os_tag} + obs = json.dumps(res) + else: + obs = json.dumps({ + "matched_count": len(matched_names), + "message": "Too many nodes matched. Showing names only.", + "node_names": matched_names + }) + except Exception as e: + obs = f"Error listing nodes: {e}" + messages.append({ + "tool_call_id": tc.id, + "role": "tool", + "name": fn, + "content": obs + }) + elif fn == "validate_playbook": + playbook_yaml = args.get("playbook_yaml", "") + validation_res = self.validate_playbook(playbook_yaml) + messages.append({ + "tool_call_id": tc.id, + "role": "tool", + "name": fn, + "content": json.dumps(validation_res) + }) + elif fn == "return_playbook": + final_playbook_yaml = args.get("playbook_yaml", "") + messages.append({ + "tool_call_id": tc.id, + "role": "tool", + "name": fn, + "content": json.dumps({"success": True, "message": "Playbook returned successfully."}) + }) + + # If return_playbook was called, we can terminate early + if final_playbook_yaml is not None: + break + + return { + "response": resp_msg.content or "", + "chat_history": messages[1:], + "playbook_yaml": final_playbook_yaml + } diff --git a/connpy/cli/ai_handler.py b/connpy/cli/ai_handler.py index 07e9e8b..c16ace0 100644 --- a/connpy/cli/ai_handler.py +++ b/connpy/cli/ai_handler.py @@ -94,7 +94,7 @@ class AIHandler: def single_question(self, args, session_id): query = " ".join(args.ask) - with console.status("[ai_status]Agent is thinking and analyzing...") as status: + with console.status("[ai_status]Agent is thinking and analyzing...[/ai_status]") as status: result = self.app.myai.ask(query, status=status, debug=args.debug, session_id=session_id, trust=args.trust, **self.ai_overrides) responder = result.get("responder", "engineer") @@ -131,7 +131,7 @@ class AIHandler: if not user_query.strip(): continue if user_query.lower() in ['exit', 'quit', 'bye', 'cancel']: break - with console.status("[ai_status]Agent is thinking...") as status: + with console.status("[ai_status]Agent is thinking...[/ai_status]") as status: result = self.app.myai.ask(user_query, chat_history=history, status=status, debug=args.debug, trust=args.trust, session_id=session_id, **self.ai_overrides) new_history = result.get("chat_history") diff --git a/connpy/cli/run_handler.py b/connpy/cli/run_handler.py index f35810f..749f63f 100644 --- a/connpy/cli/run_handler.py +++ b/connpy/cli/run_handler.py @@ -15,7 +15,12 @@ class RunHandler: def dispatch(self, args): if len(args.data) > 1: args.action = "noderun" - actions = {"noderun": self.node_run, "generate": self.yaml_generate, "run": self.yaml_run} + actions = { + "noderun": self.node_run, + "generate": self.yaml_generate, + "generate_ai": self.ai_generate, + "run": self.yaml_run + } return actions.get(args.action)(args) def node_run(self, args): @@ -33,6 +38,41 @@ class RunHandler: commands = [" ".join(args.data[1:])] + # Check for Preflight AI simulation + if getattr(args, "preflight_ai", False): + matched_node_names = [n.get("name") if isinstance(n, dict) else n for n in matched_nodes] + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Simulating execution...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[engineer][bold]Preflight AI Simulation[/bold][/engineer]", style="engineer")) + first_chunk = False + renderer.feed(chunk) + + try: + status_context.start() + self.app.services.ai.predict_execution_results( + matched_node_names, + commands, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[engineer][bold]Preflight AI Simulation[/bold][/engineer]", style="engineer")) + renderer.flush() + printer.console.print(Rule(style="engineer")) + except Exception as e: + printer.error(f"Preflight AI simulation failed: {e}") + sys.exit(1) + sys.exit(0) + try: header_printed = False @@ -70,6 +110,40 @@ class RunHandler: ) printer.run_summary(results) + # Analyze execution results if requested + if getattr(args, "analyze", None) is not None: + printer.console.print() + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Analyzing execution results...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Analysis[/bold][/architect]", style="architect")) + first_chunk = False + renderer.feed(chunk) + + query = args.analyze if args.analyze else " ".join(args.data[1:]) + try: + status_context.start() + self.app.services.ai.analyze_execution_results( + results, + query=query, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Analysis[/bold][/architect]", style="architect")) + renderer.flush() + printer.console.print(Rule(style="architect")) + except Exception as e: + printer.error(f"AI Analysis failed: {e}") + except ConnpyError as e: printer.error(str(e)) sys.exit(1) @@ -90,8 +164,105 @@ class RunHandler: with open(path, "r") as f: playbook = yaml.load(f, Loader=yaml.FullLoader) + # Check preflight first before any task runs + if getattr(args, "preflight_ai", False): + preflight_failed = False + for task in playbook.get("tasks", []): + name = task.get("name", "Task") + nodelist = task.get("nodes", []) + commands = task.get("commands", []) + + # Resolve nodes to names + try: + if isinstance(nodelist, str): + resolved_nodes = self.app.services.nodes.list_nodes(nodelist) + elif isinstance(nodelist, list): + resolved_nodes = [] + for item in nodelist: + matches = self.app.services.nodes.list_nodes(item) + for m in matches: + if m not in resolved_nodes: + resolved_nodes.append(m) + else: + resolved_nodes = [] + except Exception: + resolved_nodes = [] + + resolved_names = [n.get("name") if isinstance(n, dict) else n for n in resolved_nodes] + printer.console.print(f"\n[bold]Task: {name}[/bold] (Preflight for {len(resolved_names)} nodes)") + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Simulating execution...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title=f"[engineer][bold]Preflight AI Simulation: {name}[/bold][/engineer]", style="engineer")) + first_chunk = False + renderer.feed(chunk) + try: + status_context.start() + self.app.services.ai.predict_execution_results( + resolved_names, + commands, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title=f"[engineer][bold]Preflight AI Simulation: {name}[/bold][/engineer]", style="engineer")) + renderer.flush() + printer.console.print(Rule(style="engineer")) + except Exception as e: + printer.error(f"Preflight AI simulation failed for task {name}: {e}") + preflight_failed = True + if preflight_failed: + sys.exit(1) + sys.exit(0) + + # Standard run + results_all = {} for task in playbook.get("tasks", []): - self.cli_run(task) + task_res = self.cli_run(task) + if task_res: + results_all.update(task_res) + + # If analyze is enabled, run analysis on accumulated results + if getattr(args, "analyze", None) is not None: + printer.console.print() + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Analyzing playbook execution results...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Playbook Analysis[/bold][/architect]", style="architect")) + first_chunk = False + renderer.feed(chunk) + + query = args.analyze if args.analyze else f"Playbook: {path}" + try: + status_context.start() + self.app.services.ai.analyze_execution_results( + results_all, + query=query, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Playbook Analysis[/bold][/architect]", style="architect")) + renderer.flush() + printer.console.print(Rule(style="architect")) + except Exception as e: + printer.error(f"AI Analysis failed: {e}") except Exception as e: printer.error(f"Failed to run playbook {path}: {e}") @@ -136,6 +307,7 @@ class RunHandler: nodelist = resolved_nodes + results = {} try: header_printed = False if action == "run": @@ -195,6 +367,118 @@ class RunHandler: ) # ALWAYS show the aggregate summary at the end printer.test_summary(results) + + return results except ConnpyError as e: printer.error(str(e)) + return {} + + def ai_generate(self, args): + from rich.prompt import Prompt + from rich.rule import Rule + from rich.panel import Panel + from rich.syntax import Syntax + + dest_file = args.data[0] + if os.path.exists(dest_file): + printer.error(f"File '{dest_file}' already exists.") + sys.exit(14) + + chat_history = [] + + # Consistent layout opening matching global AI (engineer style) + from rich.markdown import Markdown + printer.console.print(Rule(style="engineer")) + printer.console.print(Markdown("**Playbook Builder AI**: Welcome! Describe the automation workflow you want to design.\nType **exit** to quit.\n")) + printer.console.print(Rule(style="engineer")) + + while True: + try: + user_prompt = Prompt.ask("[user_prompt]User[/user_prompt]") + except (KeyboardInterrupt, EOFError): + printer.console.print() + printer.warning("Operation cancelled by user.") + break + + if user_prompt.strip().lower() in ["exit", "quit"]: + printer.info("Exiting AI Assistant.") + break + + if not user_prompt.strip(): + continue + + printer.console.print() + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Agent is thinking...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: + status_context.stop() + except: + pass + printer.console.print(Rule(title="[engineer][bold]Playbook Builder AI[/bold][/engineer]", style="engineer")) + first_chunk = False + renderer.feed(chunk) + + try: + status_context.start() + res = self.app.services.ai.build_playbook_chat( + user_prompt, + chat_history=chat_history, + chunk_callback=callback + ) + if first_chunk: + try: + status_context.stop() + except: + pass + renderer.flush() + if not first_chunk: + printer.console.print(Rule(style="engineer")) + + # Update history + if res and "chat_history" in res: + chat_history = res["chat_history"] + + # Check if the agent returned a validated playbook YAML + if res and "playbook_yaml" in res and res["playbook_yaml"]: + yaml_content = res["playbook_yaml"] + printer.console.print() + printer.success("Playbook YAML successfully generated and validated.") + + # Show the YAML inside a beautiful panel matching AI style (with engineer borders) + syntax = Syntax(yaml_content, "yaml", theme="ansi_dark", word_wrap=True, background_color="default") + panel = Panel(syntax, title="[engineer][bold]Resulting Playbook[/bold][/engineer]", border_style="engineer", expand=False) + printer.console.print(panel) + + # Ask if the user wants to save it + try: + save_confirm = Prompt.ask( + f"\nDo you want to save this playbook to '{dest_file}'?", + choices=["y", "n", "run"], + default="y" + ) + except (KeyboardInterrupt, EOFError): + printer.console.print() + printer.warning("Saving skipped.") + break + + choice = save_confirm.strip().lower() + if choice in ["y", "yes", "run"]: + with open(dest_file, "w") as f: + f.write(yaml_content) + printer.success(f"Playbook saved successfully to '{dest_file}'") + if choice == "run": + printer.console.print() + printer.info("Executing the saved playbook...") + self.yaml_run(args) + break + else: + printer.warning("Playbook not saved. You can continue describing changes or exit.") + except Exception as e: + printer.error(f"Error in AI chat: {e}") diff --git a/connpy/completion.py b/connpy/completion.py index da48909..c41369c 100755 --- a/connpy/completion.py +++ b/connpy/completion.py @@ -169,12 +169,17 @@ def _build_tree(nodes, folders, profiles, plugins, configdir): run_after_node.update({ "--test": {"*": run_after_node}, "-t": {"*": run_after_node}, + "--analyze": {"*": run_after_node}, + "--preflight-ai": run_after_node, "*": run_after_node # Consume commands }) run_dict = { "--generate": {"__extra__": lambda w: get_cwd(w, "--generate")}, "-g": {"__extra__": lambda w: get_cwd(w, "-g")}, + "--generate-ai": {"__extra__": lambda w: get_cwd(w, "--generate-ai")}, + "--analyze": {"*": run_after_node}, + "--preflight-ai": run_after_node, "--test": {"*": None}, "-t": {"*": None}, "--help": None, diff --git a/connpy/connapp.py b/connpy/connapp.py index 4ebfa72..e8ff11e 100755 --- a/connpy/connapp.py +++ b/connpy/connapp.py @@ -303,6 +303,9 @@ class connapp: runparser.add_argument("run", nargs='+', action=self._store_type, help=get_help("run"), default="run").completer = nodes_completer runparser.add_argument("-t", "--test", dest="test_expected", nargs='+', help="Expected text(s) to validate in output. Converts the action from 'run' to 'test'") runparser.add_argument("-g","--generate", dest="action", action="store_const", help="Generate yaml file template", const="generate", default="run") + runparser.add_argument("--generate-ai", dest="action", action="store_const", help="Generate a playbook interactively with AI assistance", const="generate_ai") + runparser.add_argument("--analyze", nargs='?', const="", help="Analyze actual command execution results using AI") + runparser.add_argument("--preflight-ai", action="store_true", help="Simulate and predict command execution on devices using AI preventively") runparser.set_defaults(func=self._run.dispatch) #APIPARSER apiparser = subparsers.add_parser("api", help="Start and stop connpy API", description="Start and stop connpy API", formatter_class=RichHelpFormatter) diff --git a/connpy/grpc_layer/connpy_pb2.py b/connpy/grpc_layer/connpy_pb2.py index fbbaed2..3467ec8 100644 --- a/connpy/grpc_layer/connpy_pb2.py +++ b/connpy/grpc_layer/connpy_pb2.py @@ -26,7 +26,7 @@ from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0c\x63onnpy.proto\x12\x06\x63onnpy\x1a\x1cgoogle/protobuf/struct.proto\x1a\x1bgoogle/protobuf/empty.proto\"\xfc\x01\n\x0fInteractRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04sftp\x18\x02 \x01(\x08\x12\r\n\x05\x64\x65\x62ug\x18\x03 \x01(\x08\x12\x12\n\nstdin_data\x18\x04 \x01(\x0c\x12\x0c\n\x04\x63ols\x18\x05 \x01(\x05\x12\x0c\n\x04rows\x18\x06 \x01(\x05\x12\x1e\n\x16\x63onnection_params_json\x18\x07 \x01(\t\x12\x18\n\x10\x63opilot_question\x18\x08 \x01(\t\x12\x16\n\x0e\x63opilot_action\x18\t \x01(\t\x12\x1e\n\x16\x63opilot_context_buffer\x18\n \x01(\t\x12\x1e\n\x16\x63opilot_node_info_json\x18\r \x01(\t\"\x86\x02\n\x10InteractResponse\x12\x13\n\x0bstdout_data\x18\x01 \x01(\x0c\x12\x0f\n\x07success\x18\x02 \x01(\x08\x12\x15\n\rerror_message\x18\x03 \x01(\t\x12\x16\n\x0e\x63opilot_prompt\x18\x04 \x01(\x08\x12\x1e\n\x16\x63opilot_buffer_preview\x18\x05 \x01(\t\x12\x1d\n\x15\x63opilot_response_json\x18\x06 \x01(\t\x12\x1e\n\x16\x63opilot_node_info_json\x18\x07 \x01(\t\x12\x1c\n\x14\x63opilot_stream_chunk\x18\x08 \x01(\t\x12 \n\x18\x63opilot_injected_command\x18\t \x01(\t\"7\n\rFilterRequest\x12\x12\n\nfilter_str\x18\x01 \x01(\t\x12\x12\n\nformat_str\x18\x02 \x01(\t\"5\n\rValueResponse\x12$\n\x04\x64\x61ta\x18\x01 \x01(\x0b\x32\x16.google.protobuf.Value\"\x17\n\tIdRequest\x12\n\n\x02id\x18\x01 \x01(\t\"S\n\x0bNodeRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12%\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x11\n\tis_folder\x18\x03 \x01(\x08\".\n\rDeleteRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tis_folder\x18\x02 \x01(\x08\"\x1d\n\x0cMessageValue\x12\r\n\x05value\x18\x01 \x01(\t\";\n\x0bMoveRequest\x12\x0e\n\x06src_id\x18\x01 \x01(\t\x12\x0e\n\x06\x64st_id\x18\x02 \x01(\t\x12\x0c\n\x04\x63opy\x18\x03 \x01(\x08\"W\n\x0b\x42ulkRequest\x12\x0b\n\x03ids\x18\x01 \x03(\t\x12\r\n\x05hosts\x18\x02 \x03(\t\x12,\n\x0b\x63ommon_data\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\"7\n\x0eStructResponse\x12%\n\x04\x64\x61ta\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\"/\n\x0eProfileRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07resolve\x18\x02 \x01(\x08\"6\n\rStructRequest\x12%\n\x04\x64\x61ta\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\"\x1e\n\rStringRequest\x12\r\n\x05value\x18\x01 \x01(\t\"\x1f\n\x0eStringResponse\x12\r\n\x05value\x18\x01 \x01(\t\"C\n\rUpdateRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.google.protobuf.Value\"B\n\rPluginRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0bsource_file\x18\x02 \x01(\t\x12\x0e\n\x06update\x18\x03 \x01(\x08\"\xa5\x01\n\nRunRequest\x12\r\n\x05nodes\x18\x01 \x03(\t\x12\x10\n\x08\x63ommands\x18\x02 \x03(\t\x12\x0e\n\x06\x66older\x18\x03 \x01(\t\x12\x0e\n\x06prompt\x18\x04 \x01(\t\x12\x10\n\x08parallel\x18\x05 \x01(\x05\x12%\n\x04vars\x18\x06 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x0f\n\x07timeout\x18\x07 \x01(\x05\x12\x0c\n\x04name\x18\x08 \x01(\t\"\xb8\x01\n\x0bTestRequest\x12\r\n\x05nodes\x18\x01 \x03(\t\x12\x10\n\x08\x63ommands\x18\x02 \x03(\t\x12\x10\n\x08\x65xpected\x18\x03 \x03(\t\x12\x0e\n\x06\x66older\x18\x04 \x01(\t\x12\x0e\n\x06prompt\x18\x05 \x01(\t\x12\x10\n\x08parallel\x18\x06 \x01(\x05\x12%\n\x04vars\x18\x07 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x0f\n\x07timeout\x18\x08 \x01(\x05\x12\x0c\n\x04name\x18\t \x01(\t\"A\n\rScriptRequest\x12\x0e\n\x06param1\x18\x01 \x01(\t\x12\x0e\n\x06param2\x18\x02 \x01(\t\x12\x10\n\x08parallel\x18\x03 \x01(\x05\"3\n\rExportRequest\x12\x11\n\tfile_path\x18\x01 \x01(\t\x12\x0f\n\x07\x66olders\x18\x02 \x03(\t\"\x1c\n\x0bListRequest\x12\r\n\x05items\x18\x01 \x03(\t\"\x87\x03\n\nAskRequest\x12\x12\n\ninput_text\x18\x01 \x01(\t\x12\x0e\n\x06\x64ryrun\x18\x02 \x01(\x08\x12,\n\x0c\x63hat_history\x18\x03 \x01(\x0b\x32\x16.google.protobuf.Value\x12\x12\n\nsession_id\x18\x04 \x01(\t\x12\r\n\x05\x64\x65\x62ug\x18\x05 \x01(\x08\x12\x16\n\x0e\x65ngineer_model\x18\x06 \x01(\t\x12\x18\n\x10\x65ngineer_api_key\x18\x07 \x01(\t\x12\x17\n\x0f\x61rchitect_model\x18\x08 \x01(\t\x12\x19\n\x11\x61rchitect_api_key\x18\t \x01(\t\x12\r\n\x05trust\x18\n \x01(\x08\x12\x1b\n\x13\x63onfirmation_answer\x18\x0b \x01(\t\x12\x11\n\tinterrupt\x18\x0c \x01(\x08\x12.\n\rengineer_auth\x18\r \x01(\x0b\x32\x17.google.protobuf.Struct\x12/\n\x0e\x61rchitect_auth\x18\x0e \x01(\x0b\x32\x17.google.protobuf.Struct\"\xc8\x01\n\nAIResponse\x12\x12\n\ntext_chunk\x18\x01 \x01(\t\x12\x10\n\x08is_final\x18\x02 \x01(\x08\x12,\n\x0b\x66ull_result\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x15\n\rstatus_update\x18\x04 \x01(\t\x12\x15\n\rdebug_message\x18\x05 \x01(\t\x12\x1d\n\x15requires_confirmation\x18\x06 \x01(\x08\x12\x19\n\x11important_message\x18\x07 \x01(\t\"\x1d\n\x0c\x42oolResponse\x12\r\n\x05value\x18\x01 \x01(\x08\"j\n\x0fProviderRequest\x12\x10\n\x08provider\x18\x01 \x01(\t\x12\r\n\x05model\x18\x02 \x01(\t\x12\x0f\n\x07\x61pi_key\x18\x03 \x01(\t\x12%\n\x04\x61uth\x18\x04 \x01(\x0b\x32\x17.google.protobuf.Struct\"\x1b\n\nIntRequest\x12\r\n\x05value\x18\x01 \x01(\x05\"p\n\rNodeRunResult\x12\x11\n\tunique_id\x18\x01 \x01(\t\x12\x0e\n\x06output\x18\x02 \x01(\t\x12\x0e\n\x06status\x18\x03 \x01(\x05\x12,\n\x0btest_result\x18\x04 \x01(\x0b\x32\x17.google.protobuf.Struct\"m\n\x12\x46ullReplaceRequest\x12,\n\x0b\x63onnections\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\x12)\n\x08profiles\x18\x02 \x01(\x0b\x32\x17.google.protobuf.Struct\"X\n\x0e\x43opilotRequest\x12\x17\n\x0fterminal_buffer\x18\x01 \x01(\t\x12\x15\n\ruser_question\x18\x02 \x01(\t\x12\x16\n\x0enode_info_json\x18\x03 \x01(\t\"U\n\x0f\x43opilotResponse\x12\x10\n\x08\x63ommands\x18\x01 \x03(\t\x12\r\n\x05guide\x18\x02 \x01(\t\x12\x12\n\nrisk_level\x18\x03 \x01(\t\x12\r\n\x05\x65rror\x18\x04 \x01(\t\"a\n\nMCPRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0b\n\x03url\x18\x02 \x01(\t\x12\x0f\n\x07\x65nabled\x18\x03 \x01(\x08\x12\x17\n\x0f\x61uto_load_on_os\x18\x04 \x01(\t\x12\x0e\n\x06remove\x18\x05 \x01(\x08\"2\n\x0cLoginRequest\x12\x10\n\x08username\x18\x01 \x01(\t\x12\x10\n\x08password\x18\x02 \x01(\t\"D\n\rLoginResponse\x12\r\n\x05token\x18\x01 \x01(\t\x12\x10\n\x08username\x18\x02 \x01(\t\x12\x12\n\nexpires_at\x18\x03 \x01(\x03\"C\n\x15\x43hangePasswordRequest\x12\x14\n\x0cold_password\x18\x01 \x01(\t\x12\x14\n\x0cnew_password\x18\x02 \x01(\t2\xe1\x07\n\x0bNodeService\x12<\n\nlist_nodes\x12\x15.connpy.FilterRequest\x1a\x15.connpy.ValueResponse\"\x00\x12>\n\x0clist_folders\x12\x15.connpy.FilterRequest\x1a\x15.connpy.ValueResponse\"\x00\x12?\n\x10get_node_details\x12\x11.connpy.IdRequest\x1a\x16.connpy.StructResponse\"\x00\x12<\n\x0e\x65xplode_unique\x12\x11.connpy.IdRequest\x1a\x15.connpy.ValueResponse\"\x00\x12\x42\n\x0egenerate_cache\x12\x16.google.protobuf.Empty\x1a\x16.google.protobuf.Empty\"\x00\x12\x39\n\x08\x61\x64\x64_node\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\x0bupdate_node\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x12>\n\x0b\x64\x65lete_node\x12\x15.connpy.DeleteRequest\x1a\x16.google.protobuf.Empty\"\x00\x12:\n\tmove_node\x12\x13.connpy.MoveRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x39\n\x08\x62ulk_add\x12\x13.connpy.BulkRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x45\n\x16validate_parent_folder\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x12set_reserved_names\x12\x13.connpy.ListRequest\x1a\x16.google.protobuf.Empty\"\x00\x12H\n\rinteract_node\x12\x17.connpy.InteractRequest\x1a\x18.connpy.InteractResponse\"\x00(\x01\x30\x01\x12\x44\n\x0c\x66ull_replace\x12\x1a.connpy.FullReplaceRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x45\n\rget_inventory\x12\x16.google.protobuf.Empty\x1a\x1a.connpy.FullReplaceRequest\"\x00\x32\x96\x03\n\x0eProfileService\x12?\n\rlist_profiles\x12\x15.connpy.FilterRequest\x1a\x15.connpy.ValueResponse\"\x00\x12?\n\x0bget_profile\x12\x16.connpy.ProfileRequest\x1a\x16.connpy.StructResponse\"\x00\x12<\n\x0b\x61\x64\x64_profile\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x44\n\x11resolve_node_data\x12\x15.connpy.StructRequest\x1a\x16.connpy.StructResponse\"\x00\x12=\n\x0e\x64\x65lete_profile\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12?\n\x0eupdate_profile\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x32\xae\x03\n\rConfigService\x12@\n\x0cget_settings\x12\x16.google.protobuf.Empty\x1a\x16.connpy.StructResponse\"\x00\x12\x43\n\x0fget_default_dir\x12\x16.google.protobuf.Empty\x1a\x16.connpy.StringResponse\"\x00\x12\x44\n\x11set_config_folder\x12\x15.connpy.StringRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x41\n\x0eupdate_setting\x12\x15.connpy.UpdateRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x10\x65ncrypt_password\x12\x15.connpy.StringRequest\x1a\x16.connpy.StringResponse\"\x00\x12H\n\x15\x61pply_theme_from_file\x12\x15.connpy.StringRequest\x1a\x16.connpy.StructResponse\"\x00\x32\xca\x02\n\rPluginService\x12?\n\x0clist_plugins\x12\x16.google.protobuf.Empty\x1a\x15.connpy.ValueResponse\"\x00\x12=\n\nadd_plugin\x12\x15.connpy.PluginRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\rdelete_plugin\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\renable_plugin\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12=\n\x0e\x64isable_plugin\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x32\x9b\x02\n\x10\x45xecutionService\x12=\n\x0crun_commands\x12\x12.connpy.RunRequest\x1a\x15.connpy.NodeRunResult\"\x00\x30\x01\x12?\n\rtest_commands\x12\x13.connpy.TestRequest\x1a\x15.connpy.NodeRunResult\"\x00\x30\x01\x12\x41\n\x0erun_cli_script\x12\x15.connpy.ScriptRequest\x1a\x16.connpy.StructResponse\"\x00\x12\x44\n\x11run_yaml_playbook\x12\x15.connpy.ScriptRequest\x1a\x16.connpy.StructResponse\"\x00\x32\xe2\x01\n\x13ImportExportService\x12\x41\n\x0e\x65xport_to_file\x12\x15.connpy.ExportRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x10import_from_file\x12\x15.connpy.StringRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x12set_reserved_names\x12\x13.connpy.ListRequest\x1a\x16.google.protobuf.Empty\"\x00\x32\xd4\x04\n\tAIService\x12\x33\n\x03\x61sk\x12\x12.connpy.AskRequest\x1a\x12.connpy.AIResponse\"\x00(\x01\x30\x01\x12\x38\n\x07\x63onfirm\x12\x15.connpy.StringRequest\x1a\x14.connpy.BoolResponse\"\x00\x12@\n\x0b\x61sk_copilot\x12\x16.connpy.CopilotRequest\x1a\x17.connpy.CopilotResponse\"\x00\x12@\n\rlist_sessions\x12\x16.google.protobuf.Empty\x1a\x15.connpy.ValueResponse\"\x00\x12\x41\n\x0e\x64\x65lete_session\x12\x15.connpy.StringRequest\x1a\x16.google.protobuf.Empty\"\x00\x12G\n\x12\x63onfigure_provider\x12\x17.connpy.ProviderRequest\x1a\x16.google.protobuf.Empty\"\x00\x12=\n\rconfigure_mcp\x12\x12.connpy.MCPRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x10list_mcp_servers\x12\x16.google.protobuf.Empty\x1a\x15.connpy.ValueResponse\"\x00\x12\x44\n\x11load_session_data\x12\x15.connpy.StringRequest\x1a\x16.connpy.StructResponse\"\x00\x32\xc2\x02\n\rSystemService\x12\x39\n\tstart_api\x12\x12.connpy.IntRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x39\n\tdebug_api\x12\x12.connpy.IntRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\x08stop_api\x12\x16.google.protobuf.Empty\x1a\x16.google.protobuf.Empty\"\x00\x12;\n\x0brestart_api\x12\x12.connpy.IntRequest\x1a\x16.google.protobuf.Empty\"\x00\x12@\n\x0eget_api_status\x12\x16.google.protobuf.Empty\x1a\x14.connpy.BoolResponse\"\x00\x32\x91\x01\n\x0b\x41uthService\x12\x36\n\x05login\x12\x14.connpy.LoginRequest\x1a\x15.connpy.LoginResponse\"\x00\x12J\n\x0f\x63hange_password\x12\x1d.connpy.ChangePasswordRequest\x1a\x16.google.protobuf.Empty\"\x00\x62\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0c\x63onnpy.proto\x12\x06\x63onnpy\x1a\x1cgoogle/protobuf/struct.proto\x1a\x1bgoogle/protobuf/empty.proto\"\xfc\x01\n\x0fInteractRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04sftp\x18\x02 \x01(\x08\x12\r\n\x05\x64\x65\x62ug\x18\x03 \x01(\x08\x12\x12\n\nstdin_data\x18\x04 \x01(\x0c\x12\x0c\n\x04\x63ols\x18\x05 \x01(\x05\x12\x0c\n\x04rows\x18\x06 \x01(\x05\x12\x1e\n\x16\x63onnection_params_json\x18\x07 \x01(\t\x12\x18\n\x10\x63opilot_question\x18\x08 \x01(\t\x12\x16\n\x0e\x63opilot_action\x18\t \x01(\t\x12\x1e\n\x16\x63opilot_context_buffer\x18\n \x01(\t\x12\x1e\n\x16\x63opilot_node_info_json\x18\r \x01(\t\"\x86\x02\n\x10InteractResponse\x12\x13\n\x0bstdout_data\x18\x01 \x01(\x0c\x12\x0f\n\x07success\x18\x02 \x01(\x08\x12\x15\n\rerror_message\x18\x03 \x01(\t\x12\x16\n\x0e\x63opilot_prompt\x18\x04 \x01(\x08\x12\x1e\n\x16\x63opilot_buffer_preview\x18\x05 \x01(\t\x12\x1d\n\x15\x63opilot_response_json\x18\x06 \x01(\t\x12\x1e\n\x16\x63opilot_node_info_json\x18\x07 \x01(\t\x12\x1c\n\x14\x63opilot_stream_chunk\x18\x08 \x01(\t\x12 \n\x18\x63opilot_injected_command\x18\t \x01(\t\"7\n\rFilterRequest\x12\x12\n\nfilter_str\x18\x01 \x01(\t\x12\x12\n\nformat_str\x18\x02 \x01(\t\"5\n\rValueResponse\x12$\n\x04\x64\x61ta\x18\x01 \x01(\x0b\x32\x16.google.protobuf.Value\"\x17\n\tIdRequest\x12\n\n\x02id\x18\x01 \x01(\t\"S\n\x0bNodeRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12%\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x11\n\tis_folder\x18\x03 \x01(\x08\".\n\rDeleteRequest\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tis_folder\x18\x02 \x01(\x08\"\x1d\n\x0cMessageValue\x12\r\n\x05value\x18\x01 \x01(\t\";\n\x0bMoveRequest\x12\x0e\n\x06src_id\x18\x01 \x01(\t\x12\x0e\n\x06\x64st_id\x18\x02 \x01(\t\x12\x0c\n\x04\x63opy\x18\x03 \x01(\x08\"W\n\x0b\x42ulkRequest\x12\x0b\n\x03ids\x18\x01 \x03(\t\x12\r\n\x05hosts\x18\x02 \x03(\t\x12,\n\x0b\x63ommon_data\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\"7\n\x0eStructResponse\x12%\n\x04\x64\x61ta\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\"/\n\x0eProfileRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07resolve\x18\x02 \x01(\x08\"6\n\rStructRequest\x12%\n\x04\x64\x61ta\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\"\x1e\n\rStringRequest\x12\r\n\x05value\x18\x01 \x01(\t\"\x1f\n\x0eStringResponse\x12\r\n\x05value\x18\x01 \x01(\t\"C\n\rUpdateRequest\x12\x0b\n\x03key\x18\x01 \x01(\t\x12%\n\x05value\x18\x02 \x01(\x0b\x32\x16.google.protobuf.Value\"B\n\rPluginRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x13\n\x0bsource_file\x18\x02 \x01(\t\x12\x0e\n\x06update\x18\x03 \x01(\x08\"\xa5\x01\n\nRunRequest\x12\r\n\x05nodes\x18\x01 \x03(\t\x12\x10\n\x08\x63ommands\x18\x02 \x03(\t\x12\x0e\n\x06\x66older\x18\x03 \x01(\t\x12\x0e\n\x06prompt\x18\x04 \x01(\t\x12\x10\n\x08parallel\x18\x05 \x01(\x05\x12%\n\x04vars\x18\x06 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x0f\n\x07timeout\x18\x07 \x01(\x05\x12\x0c\n\x04name\x18\x08 \x01(\t\"\xb8\x01\n\x0bTestRequest\x12\r\n\x05nodes\x18\x01 \x03(\t\x12\x10\n\x08\x63ommands\x18\x02 \x03(\t\x12\x10\n\x08\x65xpected\x18\x03 \x03(\t\x12\x0e\n\x06\x66older\x18\x04 \x01(\t\x12\x0e\n\x06prompt\x18\x05 \x01(\t\x12\x10\n\x08parallel\x18\x06 \x01(\x05\x12%\n\x04vars\x18\x07 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x0f\n\x07timeout\x18\x08 \x01(\x05\x12\x0c\n\x04name\x18\t \x01(\t\"A\n\rScriptRequest\x12\x0e\n\x06param1\x18\x01 \x01(\t\x12\x0e\n\x06param2\x18\x02 \x01(\t\x12\x10\n\x08parallel\x18\x03 \x01(\x05\"3\n\rExportRequest\x12\x11\n\tfile_path\x18\x01 \x01(\t\x12\x0f\n\x07\x66olders\x18\x02 \x03(\t\"\x1c\n\x0bListRequest\x12\r\n\x05items\x18\x01 \x03(\t\"\x87\x03\n\nAskRequest\x12\x12\n\ninput_text\x18\x01 \x01(\t\x12\x0e\n\x06\x64ryrun\x18\x02 \x01(\x08\x12,\n\x0c\x63hat_history\x18\x03 \x01(\x0b\x32\x16.google.protobuf.Value\x12\x12\n\nsession_id\x18\x04 \x01(\t\x12\r\n\x05\x64\x65\x62ug\x18\x05 \x01(\x08\x12\x16\n\x0e\x65ngineer_model\x18\x06 \x01(\t\x12\x18\n\x10\x65ngineer_api_key\x18\x07 \x01(\t\x12\x17\n\x0f\x61rchitect_model\x18\x08 \x01(\t\x12\x19\n\x11\x61rchitect_api_key\x18\t \x01(\t\x12\r\n\x05trust\x18\n \x01(\x08\x12\x1b\n\x13\x63onfirmation_answer\x18\x0b \x01(\t\x12\x11\n\tinterrupt\x18\x0c \x01(\x08\x12.\n\rengineer_auth\x18\r \x01(\x0b\x32\x17.google.protobuf.Struct\x12/\n\x0e\x61rchitect_auth\x18\x0e \x01(\x0b\x32\x17.google.protobuf.Struct\"\xc8\x01\n\nAIResponse\x12\x12\n\ntext_chunk\x18\x01 \x01(\t\x12\x10\n\x08is_final\x18\x02 \x01(\x08\x12,\n\x0b\x66ull_result\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x15\n\rstatus_update\x18\x04 \x01(\t\x12\x15\n\rdebug_message\x18\x05 \x01(\t\x12\x1d\n\x15requires_confirmation\x18\x06 \x01(\x08\x12\x19\n\x11important_message\x18\x07 \x01(\t\"\x1d\n\x0c\x42oolResponse\x12\r\n\x05value\x18\x01 \x01(\x08\"j\n\x0fProviderRequest\x12\x10\n\x08provider\x18\x01 \x01(\t\x12\r\n\x05model\x18\x02 \x01(\t\x12\x0f\n\x07\x61pi_key\x18\x03 \x01(\t\x12%\n\x04\x61uth\x18\x04 \x01(\x0b\x32\x17.google.protobuf.Struct\"\x1b\n\nIntRequest\x12\r\n\x05value\x18\x01 \x01(\x05\"p\n\rNodeRunResult\x12\x11\n\tunique_id\x18\x01 \x01(\t\x12\x0e\n\x06output\x18\x02 \x01(\t\x12\x0e\n\x06status\x18\x03 \x01(\x05\x12,\n\x0btest_result\x18\x04 \x01(\x0b\x32\x17.google.protobuf.Struct\"m\n\x12\x46ullReplaceRequest\x12,\n\x0b\x63onnections\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\x12)\n\x08profiles\x18\x02 \x01(\x0b\x32\x17.google.protobuf.Struct\"X\n\x0e\x43opilotRequest\x12\x17\n\x0fterminal_buffer\x18\x01 \x01(\t\x12\x15\n\ruser_question\x18\x02 \x01(\t\x12\x16\n\x0enode_info_json\x18\x03 \x01(\t\"U\n\x0f\x43opilotResponse\x12\x10\n\x08\x63ommands\x18\x01 \x03(\t\x12\r\n\x05guide\x18\x02 \x01(\t\x12\x12\n\nrisk_level\x18\x03 \x01(\t\x12\r\n\x05\x65rror\x18\x04 \x01(\t\"a\n\nMCPRequest\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0b\n\x03url\x18\x02 \x01(\t\x12\x0f\n\x07\x65nabled\x18\x03 \x01(\x08\x12\x17\n\x0f\x61uto_load_on_os\x18\x04 \x01(\t\x12\x0e\n\x06remove\x18\x05 \x01(\x08\"2\n\x0cLoginRequest\x12\x10\n\x08username\x18\x01 \x01(\t\x12\x10\n\x08password\x18\x02 \x01(\t\"D\n\rLoginResponse\x12\r\n\x05token\x18\x01 \x01(\t\x12\x10\n\x08username\x18\x02 \x01(\t\x12\x12\n\nexpires_at\x18\x03 \x01(\x03\"C\n\x15\x43hangePasswordRequest\x12\x14\n\x0cold_password\x18\x01 \x01(\t\x12\x14\n\x0cnew_password\x18\x02 \x01(\t\"I\n\x0e\x41nalyzeRequest\x12(\n\x07results\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\r\n\x05query\x18\x02 \x01(\t\":\n\x10PreflightRequest\x12\x14\n\x0ctarget_nodes\x18\x01 \x03(\t\x12\x10\n\x08\x63ommands\x18\x02 \x03(\t2\xe1\x07\n\x0bNodeService\x12<\n\nlist_nodes\x12\x15.connpy.FilterRequest\x1a\x15.connpy.ValueResponse\"\x00\x12>\n\x0clist_folders\x12\x15.connpy.FilterRequest\x1a\x15.connpy.ValueResponse\"\x00\x12?\n\x10get_node_details\x12\x11.connpy.IdRequest\x1a\x16.connpy.StructResponse\"\x00\x12<\n\x0e\x65xplode_unique\x12\x11.connpy.IdRequest\x1a\x15.connpy.ValueResponse\"\x00\x12\x42\n\x0egenerate_cache\x12\x16.google.protobuf.Empty\x1a\x16.google.protobuf.Empty\"\x00\x12\x39\n\x08\x61\x64\x64_node\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\x0bupdate_node\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x12>\n\x0b\x64\x65lete_node\x12\x15.connpy.DeleteRequest\x1a\x16.google.protobuf.Empty\"\x00\x12:\n\tmove_node\x12\x13.connpy.MoveRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x39\n\x08\x62ulk_add\x12\x13.connpy.BulkRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x45\n\x16validate_parent_folder\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x12set_reserved_names\x12\x13.connpy.ListRequest\x1a\x16.google.protobuf.Empty\"\x00\x12H\n\rinteract_node\x12\x17.connpy.InteractRequest\x1a\x18.connpy.InteractResponse\"\x00(\x01\x30\x01\x12\x44\n\x0c\x66ull_replace\x12\x1a.connpy.FullReplaceRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x45\n\rget_inventory\x12\x16.google.protobuf.Empty\x1a\x1a.connpy.FullReplaceRequest\"\x00\x32\x96\x03\n\x0eProfileService\x12?\n\rlist_profiles\x12\x15.connpy.FilterRequest\x1a\x15.connpy.ValueResponse\"\x00\x12?\n\x0bget_profile\x12\x16.connpy.ProfileRequest\x1a\x16.connpy.StructResponse\"\x00\x12<\n\x0b\x61\x64\x64_profile\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x44\n\x11resolve_node_data\x12\x15.connpy.StructRequest\x1a\x16.connpy.StructResponse\"\x00\x12=\n\x0e\x64\x65lete_profile\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12?\n\x0eupdate_profile\x12\x13.connpy.NodeRequest\x1a\x16.google.protobuf.Empty\"\x00\x32\xae\x03\n\rConfigService\x12@\n\x0cget_settings\x12\x16.google.protobuf.Empty\x1a\x16.connpy.StructResponse\"\x00\x12\x43\n\x0fget_default_dir\x12\x16.google.protobuf.Empty\x1a\x16.connpy.StringResponse\"\x00\x12\x44\n\x11set_config_folder\x12\x15.connpy.StringRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x41\n\x0eupdate_setting\x12\x15.connpy.UpdateRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x10\x65ncrypt_password\x12\x15.connpy.StringRequest\x1a\x16.connpy.StringResponse\"\x00\x12H\n\x15\x61pply_theme_from_file\x12\x15.connpy.StringRequest\x1a\x16.connpy.StructResponse\"\x00\x32\xca\x02\n\rPluginService\x12?\n\x0clist_plugins\x12\x16.google.protobuf.Empty\x1a\x15.connpy.ValueResponse\"\x00\x12=\n\nadd_plugin\x12\x15.connpy.PluginRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\rdelete_plugin\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\renable_plugin\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x12=\n\x0e\x64isable_plugin\x12\x11.connpy.IdRequest\x1a\x16.google.protobuf.Empty\"\x00\x32\xd5\x01\n\x10\x45xecutionService\x12=\n\x0crun_commands\x12\x12.connpy.RunRequest\x1a\x15.connpy.NodeRunResult\"\x00\x30\x01\x12?\n\rtest_commands\x12\x13.connpy.TestRequest\x1a\x15.connpy.NodeRunResult\"\x00\x30\x01\x12\x41\n\x0erun_cli_script\x12\x15.connpy.ScriptRequest\x1a\x16.connpy.StructResponse\"\x00\x32\xe2\x01\n\x13ImportExportService\x12\x41\n\x0e\x65xport_to_file\x12\x15.connpy.ExportRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x10import_from_file\x12\x15.connpy.StringRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x12set_reserved_names\x12\x13.connpy.ListRequest\x1a\x16.google.protobuf.Empty\"\x00\x32\xb5\x06\n\tAIService\x12\x33\n\x03\x61sk\x12\x12.connpy.AskRequest\x1a\x12.connpy.AIResponse\"\x00(\x01\x30\x01\x12\x38\n\x07\x63onfirm\x12\x15.connpy.StringRequest\x1a\x14.connpy.BoolResponse\"\x00\x12@\n\x0b\x61sk_copilot\x12\x16.connpy.CopilotRequest\x1a\x17.connpy.CopilotResponse\"\x00\x12@\n\rlist_sessions\x12\x16.google.protobuf.Empty\x1a\x15.connpy.ValueResponse\"\x00\x12\x41\n\x0e\x64\x65lete_session\x12\x15.connpy.StringRequest\x1a\x16.google.protobuf.Empty\"\x00\x12G\n\x12\x63onfigure_provider\x12\x17.connpy.ProviderRequest\x1a\x16.google.protobuf.Empty\"\x00\x12=\n\rconfigure_mcp\x12\x12.connpy.MCPRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x43\n\x10list_mcp_servers\x12\x16.google.protobuf.Empty\x1a\x15.connpy.ValueResponse\"\x00\x12\x44\n\x11load_session_data\x12\x15.connpy.StringRequest\x1a\x16.connpy.StructResponse\"\x00\x12\x43\n\x13\x62uild_playbook_chat\x12\x12.connpy.AskRequest\x1a\x12.connpy.AIResponse\"\x00(\x01\x30\x01\x12K\n\x19\x61nalyze_execution_results\x12\x16.connpy.AnalyzeRequest\x1a\x12.connpy.AIResponse\"\x00\x30\x01\x12M\n\x19predict_execution_results\x12\x18.connpy.PreflightRequest\x1a\x12.connpy.AIResponse\"\x00\x30\x01\x32\xc2\x02\n\rSystemService\x12\x39\n\tstart_api\x12\x12.connpy.IntRequest\x1a\x16.google.protobuf.Empty\"\x00\x12\x39\n\tdebug_api\x12\x12.connpy.IntRequest\x1a\x16.google.protobuf.Empty\"\x00\x12<\n\x08stop_api\x12\x16.google.protobuf.Empty\x1a\x16.google.protobuf.Empty\"\x00\x12;\n\x0brestart_api\x12\x12.connpy.IntRequest\x1a\x16.google.protobuf.Empty\"\x00\x12@\n\x0eget_api_status\x12\x16.google.protobuf.Empty\x1a\x14.connpy.BoolResponse\"\x00\x32\x91\x01\n\x0b\x41uthService\x12\x36\n\x05login\x12\x14.connpy.LoginRequest\x1a\x15.connpy.LoginResponse\"\x00\x12J\n\x0f\x63hange_password\x12\x1d.connpy.ChangePasswordRequest\x1a\x16.google.protobuf.Empty\"\x00\x62\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -103,22 +103,26 @@ if not _descriptor._USE_C_DESCRIPTORS: _globals['_LOGINRESPONSE']._serialized_end=3309 _globals['_CHANGEPASSWORDREQUEST']._serialized_start=3311 _globals['_CHANGEPASSWORDREQUEST']._serialized_end=3378 - _globals['_NODESERVICE']._serialized_start=3381 - _globals['_NODESERVICE']._serialized_end=4374 - _globals['_PROFILESERVICE']._serialized_start=4377 - _globals['_PROFILESERVICE']._serialized_end=4783 - _globals['_CONFIGSERVICE']._serialized_start=4786 - _globals['_CONFIGSERVICE']._serialized_end=5216 - _globals['_PLUGINSERVICE']._serialized_start=5219 - _globals['_PLUGINSERVICE']._serialized_end=5549 - _globals['_EXECUTIONSERVICE']._serialized_start=5552 - _globals['_EXECUTIONSERVICE']._serialized_end=5835 - _globals['_IMPORTEXPORTSERVICE']._serialized_start=5838 - _globals['_IMPORTEXPORTSERVICE']._serialized_end=6064 - _globals['_AISERVICE']._serialized_start=6067 - _globals['_AISERVICE']._serialized_end=6663 - _globals['_SYSTEMSERVICE']._serialized_start=6666 - _globals['_SYSTEMSERVICE']._serialized_end=6988 - _globals['_AUTHSERVICE']._serialized_start=6991 - _globals['_AUTHSERVICE']._serialized_end=7136 + _globals['_ANALYZEREQUEST']._serialized_start=3380 + _globals['_ANALYZEREQUEST']._serialized_end=3453 + _globals['_PREFLIGHTREQUEST']._serialized_start=3455 + _globals['_PREFLIGHTREQUEST']._serialized_end=3513 + _globals['_NODESERVICE']._serialized_start=3516 + _globals['_NODESERVICE']._serialized_end=4509 + _globals['_PROFILESERVICE']._serialized_start=4512 + _globals['_PROFILESERVICE']._serialized_end=4918 + _globals['_CONFIGSERVICE']._serialized_start=4921 + _globals['_CONFIGSERVICE']._serialized_end=5351 + _globals['_PLUGINSERVICE']._serialized_start=5354 + _globals['_PLUGINSERVICE']._serialized_end=5684 + _globals['_EXECUTIONSERVICE']._serialized_start=5687 + _globals['_EXECUTIONSERVICE']._serialized_end=5900 + _globals['_IMPORTEXPORTSERVICE']._serialized_start=5903 + _globals['_IMPORTEXPORTSERVICE']._serialized_end=6129 + _globals['_AISERVICE']._serialized_start=6132 + _globals['_AISERVICE']._serialized_end=6953 + _globals['_SYSTEMSERVICE']._serialized_start=6956 + _globals['_SYSTEMSERVICE']._serialized_end=7278 + _globals['_AUTHSERVICE']._serialized_start=7281 + _globals['_AUTHSERVICE']._serialized_end=7426 # @@protoc_insertion_point(module_scope) diff --git a/connpy/grpc_layer/connpy_pb2_grpc.py b/connpy/grpc_layer/connpy_pb2_grpc.py index cb8e5b7..28f077c 100644 --- a/connpy/grpc_layer/connpy_pb2_grpc.py +++ b/connpy/grpc_layer/connpy_pb2_grpc.py @@ -1542,11 +1542,6 @@ class ExecutionServiceStub(object): request_serializer=connpy__pb2.ScriptRequest.SerializeToString, response_deserializer=connpy__pb2.StructResponse.FromString, _registered_method=True) - self.run_yaml_playbook = channel.unary_unary( - '/connpy.ExecutionService/run_yaml_playbook', - request_serializer=connpy__pb2.ScriptRequest.SerializeToString, - response_deserializer=connpy__pb2.StructResponse.FromString, - _registered_method=True) class ExecutionServiceServicer(object): @@ -1570,12 +1565,6 @@ class ExecutionServiceServicer(object): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') - def run_yaml_playbook(self, request, context): - """Missing associated documentation comment in .proto file.""" - context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details('Method not implemented!') - raise NotImplementedError('Method not implemented!') - def add_ExecutionServiceServicer_to_server(servicer, server): rpc_method_handlers = { @@ -1594,11 +1583,6 @@ def add_ExecutionServiceServicer_to_server(servicer, server): request_deserializer=connpy__pb2.ScriptRequest.FromString, response_serializer=connpy__pb2.StructResponse.SerializeToString, ), - 'run_yaml_playbook': grpc.unary_unary_rpc_method_handler( - servicer.run_yaml_playbook, - request_deserializer=connpy__pb2.ScriptRequest.FromString, - response_serializer=connpy__pb2.StructResponse.SerializeToString, - ), } generic_handler = grpc.method_handlers_generic_handler( 'connpy.ExecutionService', rpc_method_handlers) @@ -1691,33 +1675,6 @@ class ExecutionService(object): metadata, _registered_method=True) - @staticmethod - def run_yaml_playbook(request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None): - return grpc.experimental.unary_unary( - request, - target, - '/connpy.ExecutionService/run_yaml_playbook', - connpy__pb2.ScriptRequest.SerializeToString, - connpy__pb2.StructResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - _registered_method=True) - class ImportExportServiceStub(object): """Missing associated documentation comment in .proto file.""" @@ -1931,6 +1888,21 @@ class AIServiceStub(object): request_serializer=connpy__pb2.StringRequest.SerializeToString, response_deserializer=connpy__pb2.StructResponse.FromString, _registered_method=True) + self.build_playbook_chat = channel.stream_stream( + '/connpy.AIService/build_playbook_chat', + request_serializer=connpy__pb2.AskRequest.SerializeToString, + response_deserializer=connpy__pb2.AIResponse.FromString, + _registered_method=True) + self.analyze_execution_results = channel.unary_stream( + '/connpy.AIService/analyze_execution_results', + request_serializer=connpy__pb2.AnalyzeRequest.SerializeToString, + response_deserializer=connpy__pb2.AIResponse.FromString, + _registered_method=True) + self.predict_execution_results = channel.unary_stream( + '/connpy.AIService/predict_execution_results', + request_serializer=connpy__pb2.PreflightRequest.SerializeToString, + response_deserializer=connpy__pb2.AIResponse.FromString, + _registered_method=True) class AIServiceServicer(object): @@ -1990,6 +1962,24 @@ class AIServiceServicer(object): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') + def build_playbook_chat(self, request_iterator, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def analyze_execution_results(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def predict_execution_results(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + def add_AIServiceServicer_to_server(servicer, server): rpc_method_handlers = { @@ -2038,6 +2028,21 @@ def add_AIServiceServicer_to_server(servicer, server): request_deserializer=connpy__pb2.StringRequest.FromString, response_serializer=connpy__pb2.StructResponse.SerializeToString, ), + 'build_playbook_chat': grpc.stream_stream_rpc_method_handler( + servicer.build_playbook_chat, + request_deserializer=connpy__pb2.AskRequest.FromString, + response_serializer=connpy__pb2.AIResponse.SerializeToString, + ), + 'analyze_execution_results': grpc.unary_stream_rpc_method_handler( + servicer.analyze_execution_results, + request_deserializer=connpy__pb2.AnalyzeRequest.FromString, + response_serializer=connpy__pb2.AIResponse.SerializeToString, + ), + 'predict_execution_results': grpc.unary_stream_rpc_method_handler( + servicer.predict_execution_results, + request_deserializer=connpy__pb2.PreflightRequest.FromString, + response_serializer=connpy__pb2.AIResponse.SerializeToString, + ), } generic_handler = grpc.method_handlers_generic_handler( 'connpy.AIService', rpc_method_handlers) @@ -2292,6 +2297,87 @@ class AIService(object): metadata, _registered_method=True) + @staticmethod + def build_playbook_chat(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_stream( + request_iterator, + target, + '/connpy.AIService/build_playbook_chat', + connpy__pb2.AskRequest.SerializeToString, + connpy__pb2.AIResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def analyze_execution_results(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/connpy.AIService/analyze_execution_results', + connpy__pb2.AnalyzeRequest.SerializeToString, + connpy__pb2.AIResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def predict_execution_results(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/connpy.AIService/predict_execution_results', + connpy__pb2.PreflightRequest.SerializeToString, + connpy__pb2.AIResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + class SystemServiceStub(object): """Missing associated documentation comment in .proto file.""" diff --git a/connpy/grpc_layer/server.py b/connpy/grpc_layer/server.py index 1919990..03fa2b9 100644 --- a/connpy/grpc_layer/server.py +++ b/connpy/grpc_layer/server.py @@ -791,11 +791,6 @@ class ExecutionServicer(connpy_pb2_grpc.ExecutionServiceServicer): res = self.service.run_cli_script(request.param1, request.param2, request.parallel) return connpy_pb2.StructResponse(data=to_struct(res)) - @handle_errors - def run_yaml_playbook(self, request, context): - res = self.service.run_yaml_playbook(request.param1, request.parallel) - return connpy_pb2.StructResponse(data=to_struct(res)) - class ImportExportServicer(connpy_pb2_grpc.ImportExportServiceServicer): def __init__(self, provider, registry=None): if not hasattr(provider, "mode"): @@ -955,12 +950,10 @@ class AIServicer(connpy_pb2_grpc.AIServiceServicer): def service(self): return self._get_provider().ai - @handle_errors - def ask(self, request_iterator, context): + def _handle_chat_stream(self, request_iterator, context, service_method): import queue import threading - ai_service = self.service chunk_queue = queue.Queue() request_queue = queue.Queue() bridge = None @@ -978,21 +971,28 @@ class AIServicer(connpy_pb2_grpc.AIServiceServicer): nonlocal history, bridge, agent_instance try: # Run the AI interaction (this blocks this specific thread) - res = ai_service.ask( - input_text, - chat_history=history if history else None, - session_id=session_id, - debug=debug, - status=bridge, - console=bridge, - confirm_handler=bridge.confirm, - chunk_callback=callback, - trust=trust, - **overrides - ) + if getattr(service_method, "__name__", None) == "build_playbook_chat": + res = service_method( + input_text, + chat_history=history if history else None, + status=bridge, + chunk_callback=callback + ) + else: + res = service_method( + input_text, + chat_history=history if history else None, + session_id=session_id, + debug=debug, + status=bridge, + confirm_handler=bridge.confirm, + chunk_callback=callback, + trust=trust, + **overrides + ) # Update history for next message - if "chat_history" in res: + if res and "chat_history" in res: history = res["chat_history"] # Send final chunk marker @@ -1086,6 +1086,71 @@ class AIServicer(connpy_pb2_grpc.AIServiceServicer): elif msg_type == "final_mark": yield connpy_pb2.AIResponse(is_final=True, full_result=to_struct(val)) + def _handle_unary_stream(self, service_method, *args, **kwargs): + import queue + import threading + + chunk_queue = queue.Queue() + bridge = StatusBridge(chunk_queue, is_web=False) + + def callback(chunk): + chunk_queue.put(("text", chunk)) + + def _worker(): + try: + res = service_method(*args, chunk_callback=callback, status=bridge, **kwargs) + chunk_queue.put(("final_mark", res)) + except Exception as e: + import traceback + print(f"gRPC Unary Stream error: {e}") + traceback.print_exc() + chunk_queue.put(("status", f"Error: {str(e)}")) + chunk_queue.put(("final_mark", {"response": f"Error: {str(e)}", "error": True})) + finally: + chunk_queue.put((None, None)) + + threading.Thread(target=_worker, daemon=True).start() + + while True: + item = chunk_queue.get() + if item == (None, None): + break + + msg_type, val = item + if msg_type == "text": + yield connpy_pb2.AIResponse(text_chunk=val, is_final=False) + elif msg_type == "status": + clean_val = val.replace("[ai_status]", "").replace("[/ai_status]", "") + yield connpy_pb2.AIResponse(status_update=clean_val, is_final=False) + elif msg_type == "debug": + yield connpy_pb2.AIResponse(debug_message=val, is_final=False) + elif msg_type == "important": + yield connpy_pb2.AIResponse(important_message=val, is_final=False) + elif msg_type == "confirm": + yield connpy_pb2.AIResponse(status_update=val, requires_confirmation=True, is_final=False) + elif msg_type == "final_mark": + yield connpy_pb2.AIResponse(is_final=True, full_result=to_struct(val)) + + @handle_errors + def ask(self, request_iterator, context): + yield from self._handle_chat_stream(request_iterator, context, self.service.ask) + + @handle_errors + def build_playbook_chat(self, request_iterator, context): + yield from self._handle_chat_stream(request_iterator, context, self.service.build_playbook_chat) + + @handle_errors + def analyze_execution_results(self, request, context): + results = from_struct(request.results) + query = request.query if request.query else None + yield from self._handle_unary_stream(self.service.analyze_execution_results, results, query=query) + + @handle_errors + def predict_execution_results(self, request, context): + target_nodes = list(request.target_nodes) + commands = list(request.commands) + yield from self._handle_unary_stream(self.service.predict_execution_results, target_nodes, commands) + @handle_errors def confirm(self, request, context): res = self.service.confirm(request.value) diff --git a/connpy/grpc_layer/stubs.py b/connpy/grpc_layer/stubs.py index 56acd38..3971af8 100644 --- a/connpy/grpc_layer/stubs.py +++ b/connpy/grpc_layer/stubs.py @@ -692,11 +692,6 @@ class ExecutionStub: req = connpy_pb2.ScriptRequest(param1=nodes_filter, param2=script_path, parallel=parallel) return from_struct(self.stub.run_cli_script(req).data) - @handle_errors - def run_yaml_playbook(self, playbook_path, parallel=10): - req = connpy_pb2.ScriptRequest(param1=playbook_path, parallel=parallel) - return from_struct(self.stub.run_yaml_playbook(req).data) - class ImportExportStub: def __init__(self, channel, remote_host): self.stub = connpy_pb2_grpc.ImportExportServiceStub(channel) @@ -724,8 +719,7 @@ class AIStub: self.stub = connpy_pb2_grpc.AIServiceStub(channel) self.remote_host = remote_host - @handle_errors - def ask(self, input_text, dryrun=False, chat_history=None, session_id=None, debug=False, status=None, **overrides): + def _ai_chat_stream(self, stub_method, input_text, dryrun=False, chat_history=None, session_id=None, debug=False, status=None, chunk_callback=None, **overrides): import queue from rich.prompt import Prompt from rich.text import Text @@ -760,7 +754,7 @@ class AIStub: if req is None: break yield req - responses = self.stub.ask(request_generator()) + responses = stub_method(request_generator()) full_content = "" header_printed = False @@ -859,26 +853,32 @@ class AIStub: try: status.stop() except: pass - from rich.console import Console as RichConsole - from rich.rule import Rule - from ..printer import connpy_theme, get_original_stdout, IncrementalMarkdownParser - stable_console = RichConsole(theme=connpy_theme, file=get_original_stdout()) - - # Print header on first chunk - alias = "architect" if current_responder == "architect" else "engineer" - role_label = "Network Architect" if current_responder == "architect" else "Network Engineer" - stable_console.print(Rule(f"[bold {alias}]{role_label}[/bold {alias}]", style=alias)) - header_printed = True - - # Initialize parser - md_parser = IncrementalMarkdownParser(console=stable_console) + if chunk_callback: + header_printed = True + else: + from rich.console import Console as RichConsole + from rich.rule import Rule + from ..printer import connpy_theme, get_original_stdout, IncrementalMarkdownParser + stable_console = RichConsole(theme=connpy_theme, file=get_original_stdout()) + + # Print header on first chunk + alias = "architect" if current_responder == "architect" else "engineer" + role_label = "Network Architect" if current_responder == "architect" else "Network Engineer" + stable_console.print(Rule(f"[bold {alias}]{role_label}[/bold {alias}]", style=alias)) + header_printed = True + + # Initialize parser + md_parser = IncrementalMarkdownParser(console=stable_console) full_content += response.text_chunk - md_parser.feed(response.text_chunk) + if chunk_callback: + chunk_callback(response.text_chunk) + elif md_parser: + md_parser.feed(response.text_chunk) continue if response.is_final: - if header_printed: + if not chunk_callback and header_printed: from rich.rule import Rule md_parser.flush() @@ -887,12 +887,8 @@ class AIStub: except: pass final_result = from_struct(response.full_result) - responder = final_result.get("responder", "engineer") - alias = "architect" if responder == "architect" else "engineer" - role_label = "Network Architect" if responder == "architect" else "Network Engineer" - title = f"[bold {alias}]{role_label}[/bold {alias}]" - if header_printed: + if not chunk_callback and header_printed: from rich.console import Console as RichConsole from ..printer import connpy_theme, get_original_stdout stable_console = RichConsole(theme=connpy_theme, file=get_original_stdout()) @@ -911,6 +907,104 @@ class AIStub: return final_result + @handle_errors + def ask(self, input_text, dryrun=False, chat_history=None, session_id=None, debug=False, status=None, **overrides): + return self._ai_chat_stream(self.stub.ask, input_text, dryrun=dryrun, chat_history=chat_history, session_id=session_id, debug=debug, status=status, **overrides) + + @handle_errors + def build_playbook_chat(self, user_input, chat_history=None, status=None, chunk_callback=None): + return self._ai_chat_stream(self.stub.build_playbook_chat, user_input, chat_history=chat_history, status=status, chunk_callback=chunk_callback) + + def _process_unary_stream(self, responses, status=None, chunk_callback=None): + full_content = "" + header_printed = False + final_result = {"response": "", "chat_history": []} + md_parser = None + + try: + for response in responses: + if response.status_update: + if status: + status.update(response.status_update) + continue + + if response.important_message: + if status: + try: status.stop() + except: pass + printer.console.print(Text.from_ansi(response.important_message)) + if status: + try: status.start() + except: pass + continue + + if not response.is_final: + if response.text_chunk: + if not header_printed: + if status: + try: status.stop() + except: pass + + if chunk_callback: + header_printed = True + else: + from rich.console import Console as RichConsole + from rich.rule import Rule + from ..printer import connpy_theme, get_original_stdout, IncrementalMarkdownParser + stable_console = RichConsole(theme=connpy_theme, file=get_original_stdout()) + + # Print default header + stable_console.print(Rule("[bold engineer]AI Analysis[/bold engineer]", style="engineer")) + header_printed = True + md_parser = IncrementalMarkdownParser(console=stable_console) + + full_content += response.text_chunk + if chunk_callback: + chunk_callback(response.text_chunk) + elif md_parser: + md_parser.feed(response.text_chunk) + continue + + if response.is_final: + if md_parser: + md_parser.flush() + + if status: + try: status.stop() + except: pass + + final_result = from_struct(response.full_result) + + if md_parser: + from rich.console import Console as RichConsole + from rich.rule import Rule + from ..printer import connpy_theme, get_original_stdout + stable_console = RichConsole(theme=connpy_theme, file=get_original_stdout()) + stable_console.print(Rule(style="engineer")) + break + except Exception as e: + if isinstance(e, grpc.RpcError): + raise + printer.warning(f"Stream interrupted: {e}") + + if full_content: + final_result["streamed"] = True + + return final_result + + @handle_errors + def analyze_execution_results(self, results, query=None, status=None, chunk_callback=None): + req = connpy_pb2.AnalyzeRequest(query=query or "") + req.results.CopyFrom(to_struct(results)) + responses = self.stub.analyze_execution_results(req) + return self._process_unary_stream(responses, status, chunk_callback) + + @handle_errors + def predict_execution_results(self, target_nodes, commands, status=None, chunk_callback=None): + req = connpy_pb2.PreflightRequest(target_nodes=target_nodes, commands=commands) + responses = self.stub.predict_execution_results(req) + return self._process_unary_stream(responses, status, chunk_callback) + @handle_errors def confirm(self, input_text, console=None): return self.stub.confirm(connpy_pb2.StringRequest(value=input_text)).value diff --git a/connpy/printer.py b/connpy/printer.py index 6cbd071..cf885b0 100644 --- a/connpy/printer.py +++ b/connpy/printer.py @@ -573,7 +573,7 @@ class BlockMarkdownRenderer: if not block_text: return from rich.markdown import Markdown - self._console.print(Markdown(block_text)) + self._console.print(Markdown(block_text, code_theme="ansi_dark")) # Alias for backward compatibility IncrementalMarkdownParser = BlockMarkdownRenderer diff --git a/connpy/proto/connpy.proto b/connpy/proto/connpy.proto index c3c953d..842ff73 100644 --- a/connpy/proto/connpy.proto +++ b/connpy/proto/connpy.proto @@ -53,7 +53,6 @@ service ExecutionService { rpc run_commands (RunRequest) returns (stream NodeRunResult) {} rpc test_commands (TestRequest) returns (stream NodeRunResult) {} rpc run_cli_script (ScriptRequest) returns (StructResponse) {} - rpc run_yaml_playbook (ScriptRequest) returns (StructResponse) {} } service ImportExportService { @@ -72,6 +71,9 @@ service AIService { rpc configure_mcp (MCPRequest) returns (google.protobuf.Empty) {} rpc list_mcp_servers (google.protobuf.Empty) returns (ValueResponse) {} rpc load_session_data (StringRequest) returns (StructResponse) {} + rpc build_playbook_chat (stream AskRequest) returns (stream AIResponse) {} + rpc analyze_execution_results (AnalyzeRequest) returns (stream AIResponse) {} + rpc predict_execution_results (PreflightRequest) returns (stream AIResponse) {} } service SystemService { @@ -317,3 +319,13 @@ message ChangePasswordRequest { string old_password = 1; string new_password = 2; } + +message AnalyzeRequest { + google.protobuf.Struct results = 1; + string query = 2; +} + +message PreflightRequest { + repeated string target_nodes = 1; + repeated string commands = 2; +} diff --git a/connpy/services/ai_service.py b/connpy/services/ai_service.py index 4548154..dc81d81 100644 --- a/connpy/services/ai_service.py +++ b/connpy/services/ai_service.py @@ -319,3 +319,37 @@ class AIService(BaseService): agent = ai(self.config) return agent.load_session_data(session_id) + def build_playbook_chat(self, user_input: str, chat_history: list = None, status=None, chunk_callback=None): + """Interact with the specialized Playbook Builder Agent.""" + from connpy.ai import PlaybookBuilderAgent + agent = PlaybookBuilderAgent(self.config) + return agent.ask(user_input, chat_history=chat_history, status=status, chunk_callback=chunk_callback) + + def analyze_execution_results(self, results: dict, query: str = None, status=None, chunk_callback=None): + """Analyze actual command execution results using Network Architect 1-shot.""" + import json + results_str = json.dumps(results, indent=2) + + prompt = f"@architect: Please analyze the following actual execution results. Diagnose any issues, highlight successful actions, and suggest strategic remediation steps if needed." + if query: + prompt += f"\nSpecific user request: {query}" + prompt += f"\n\nResults Data:\n{results_str}" + prompt += "\n\nCRITICAL DIRECTIVE: You are running in a strictly 1-shot offline diagnostics mode (--analyze). There is no active conversation loop, and you are NOT conversing with a Network Engineer. You MUST deliver your complete strategic analysis immediately. DO NOT suggest, mention, or attempt to delegate the session back to the engineer." + + # Delegate to self.ask, setting stream=True and forwarding callback/status. + # This will invoke standard ai.ask with '@architect:' prefix, forcing 1-shot architect brain. + return self.ask(prompt, status=status, chunk_callback=chunk_callback, one_shot=True) + + def predict_execution_results(self, target_nodes: list, commands: list, status=None, chunk_callback=None): + """Predict and simulate execution results preventively using the Preflight Simulation Agent (1-shot).""" + nodes_str = ", ".join(target_nodes) + commands_str = "\n".join(f"- {cmd}" for cmd in commands) + + prompt = f"@engineer: Act as a Preflight Simulation Agent. Simulate and predict the expected outputs and behaviors of the following commands on the target nodes. Alert about potential safety or configuration risks based on node profiles." + prompt += f"\n\nTarget Nodes: {nodes_str}" + prompt += f"\nCommands to simulate:\n{commands_str}" + prompt += "\n\nCRITICAL SCALABILITY DIRECTIVE: If there are many target nodes, DO NOT list predictions node-by-node. Instead, group them by Operating System, vendor, or platform, and provide a highly concise Executive Summary. Detail individual risks only for nodes that present specific anomalies or security concerns. Focus on overall impact." + + # Delegate to self.ask, using the standard engineer brain but with the simulated preflight prompt. + return self.ask(prompt, status=status, chunk_callback=chunk_callback) + diff --git a/connpy/services/execution_service.py b/connpy/services/execution_service.py index 7588e16..ae58f89 100644 --- a/connpy/services/execution_service.py +++ b/connpy/services/execution_service.py @@ -1,6 +1,5 @@ from typing import List, Dict, Any, Callable, Optional import os -import yaml from .base import BaseService from connpy.core import nodes as Nodes from .exceptions import ConnpyError @@ -108,52 +107,3 @@ class ExecutionService(BaseService): return self.run_commands(nodes_filter, commands, parallel=parallel) - def run_yaml_playbook(self, playbook_data: str, parallel: int = 10) -> Dict[str, Any]: - """Run a structured Connpy YAML automation playbook (from path or content).""" - playbook = None - if playbook_data.startswith("---YAML---\n"): - try: - content = playbook_data[len("---YAML---\n"):] - playbook = yaml.load(content, Loader=yaml.FullLoader) - except Exception as e: - raise ConnpyError(f"Failed to parse YAML content: {e}") - else: - if not os.path.exists(playbook_data): - raise ConnpyError(f"Playbook file not found: {playbook_data}") - try: - with open(playbook_data, "r") as f: - playbook = yaml.load(f, Loader=yaml.FullLoader) - except Exception as e: - raise ConnpyError(f"Failed to load playbook {playbook_data}: {e}") - - # Basic validation - if not isinstance(playbook, dict) or "nodes" not in playbook or "commands" not in playbook: - raise ConnpyError("Invalid playbook format: missing 'nodes' or 'commands' keys.") - - action = playbook.get("action", "run") - options = playbook.get("options", {}) - - # Extract all fields similar to RunHandler.cli_run - exec_args = { - "nodes_filter": playbook["nodes"], - "commands": playbook["commands"], - "variables": playbook.get("variables"), - "parallel": options.get("parallel", parallel), - "timeout": playbook.get("timeout", options.get("timeout", 20)), - "prompt": options.get("prompt"), - "name": playbook.get("name", "Task") - } - - # Map 'output' field to folder path if it's not stdout/null - output_cfg = playbook.get("output") - if output_cfg not in [None, "stdout"]: - exec_args["folder"] = output_cfg - - if action == "run": - return self.run_commands(**exec_args) - elif action == "test": - exec_args["expected"] = playbook.get("expected", []) - return self.test_commands(**exec_args) - else: - raise ConnpyError(f"Unsupported playbook action: {action}") - diff --git a/connpy/tests/test_ai.py b/connpy/tests/test_ai.py index b0e4bcf..d878920 100644 --- a/connpy/tests/test_ai.py +++ b/connpy/tests/test_ai.py @@ -480,6 +480,15 @@ class TestToolDefinitions: names = [t["function"]["name"] for t in tools] assert "arch_tool" in names + def test_architect_tools_one_shot(self, ai_config): + from connpy.ai import ai + one_shot_ai = ai(ai_config, one_shot=True) + tools = one_shot_ai._get_architect_tools() + names = [t["function"]["name"] for t in tools] + assert "delegate_to_engineer" not in names + assert "return_to_engineer" not in names + assert "manage_memory_tool" in names + # ========================================================================= # AI Session Management tests diff --git a/connpy/tests/test_cli_run_ai.py b/connpy/tests/test_cli_run_ai.py new file mode 100644 index 0000000..50bfbb0 --- /dev/null +++ b/connpy/tests/test_cli_run_ai.py @@ -0,0 +1,136 @@ +import pytest +from unittest.mock import patch, MagicMock, ANY +from connpy.connapp import connapp +import os + +@pytest.fixture +def app(populated_config): + """Returns an instance of connapp initialized with mock config.""" + return connapp(populated_config) + +def test_run_generate_ai_dispatch(app): + """Test that connpy run --generate-ai parses and calls ai_generate.""" + with patch("connpy.cli.run_handler.RunHandler.ai_generate") as mock_ai_gen: + app.start(["run", "--generate-ai", "new_playbook.yaml"]) + mock_ai_gen.assert_called_once() + args = mock_ai_gen.call_args[0][0] + assert args.data == ["new_playbook.yaml"] + assert args.action == "generate_ai" + +def test_run_preflight_ai_node(app): + """Test that connpy run --preflight-ai calls predict_execution_results and exits.""" + with patch("connpy.services.node_service.NodeService.list_nodes", return_value=["router1"]): + with patch("connpy.services.ai_service.AIService.predict_execution_results") as mock_predict: + with pytest.raises(SystemExit) as exc: + app.start(["run", "router1", "show version", "--preflight-ai"]) + + assert exc.value.code == 0 + mock_predict.assert_called_once_with(["router1"], ["show version"], chunk_callback=ANY) + +def test_run_analyze_node(app): + """Test that connpy run --analyze calls analyze_execution_results after execution.""" + mock_run = MagicMock(return_value={"router1": {"status": 0, "output": "success"}}) + + with patch("connpy.services.node_service.NodeService.list_nodes", return_value=["router1"]): + with patch("connpy.services.execution_service.ExecutionService.run_commands", mock_run): + with patch("connpy.services.ai_service.AIService.analyze_execution_results") as mock_analyze: + app.start(["run", "router1", "show version", "--analyze"]) + mock_run.assert_called_once() + mock_analyze.assert_called_once_with( + {"router1": {"status": 0, "output": "success"}}, + query="show version", + chunk_callback=ANY + ) + +def test_run_preflight_ai_playbook(app, tmp_path): + """Test that running a playbook with --preflight-ai predicts results per task.""" + playbook_path = tmp_path / "test_playbook.yaml" + playbook_content = """ +tasks: + - name: test-task + action: run + nodes: "router1" + commands: ["show ip interface brief"] + output: stdout +""" + playbook_path.write_text(playbook_content) + + with patch("connpy.services.node_service.NodeService.list_nodes", return_value=["router1"]): + with patch("connpy.services.ai_service.AIService.predict_execution_results") as mock_predict: + with pytest.raises(SystemExit) as exc: + app.start(["run", str(playbook_path), "--preflight-ai"]) + + assert exc.value.code == 0 + mock_predict.assert_called_once_with(["router1"], ["show ip interface brief"], chunk_callback=ANY) + +def test_run_analyze_playbook(app, tmp_path): + """Test that running a playbook with --analyze triggers strategic analysis on all task outcomes.""" + playbook_path = tmp_path / "test_playbook.yaml" + playbook_content = """ +tasks: + - name: test-task + action: run + nodes: "router1" + commands: ["show ip interface brief"] + output: stdout +""" + playbook_path.write_text(playbook_content) + + mock_run = MagicMock(return_value={"router1": {"status": 0, "output": "ok"}}) + + with patch("connpy.services.node_service.NodeService.list_nodes", return_value=["router1"]): + with patch("connpy.services.execution_service.ExecutionService.run_commands", mock_run): + with patch("connpy.services.ai_service.AIService.analyze_execution_results") as mock_analyze: + app.start(["run", str(playbook_path), "--analyze"]) + mock_run.assert_called_once() + mock_analyze.assert_called_once_with( + {"router1": {"status": 0, "output": "ok"}}, + query=f"Playbook: {str(playbook_path)}", + chunk_callback=ANY + ) + +def test_ai_generate_wizard_save(app, tmp_path): + """Test that ai_generate wizard runs interactive chat loop, asks for validation and saves YAML.""" + dest_yaml = tmp_path / "playbook.yaml" + + mock_chat = MagicMock(return_value={ + "response": "Here is your playbook.", + "chat_history": [], + "playbook_yaml": "tasks:\n - name: mytask" + }) + app.services.ai.build_playbook_chat = mock_chat + + # Mock rich.prompt.Prompt.ask to simulate User inputting prompt and then 'y' to save + with patch("rich.prompt.Prompt.ask", side_effect=["create a basic task", "y"]): + app.start(["run", "--generate-ai", str(dest_yaml)]) + + mock_chat.assert_called_once_with("create a basic task", chat_history=[], chunk_callback=ANY) + assert os.path.exists(dest_yaml) + with open(dest_yaml) as f: + content = f.read() + assert "tasks:" in content + +def test_ai_generate_wizard_run(app, tmp_path): + """Test that ai_generate wizard runs, saves the playbook and executes it when choosing 'run'.""" + dest_yaml = tmp_path / "playbook_run.yaml" + + mock_chat = MagicMock(return_value={ + "response": "Here is your playbook.", + "chat_history": [], + "playbook_yaml": "tasks:\n - name: mytask\n action: run\n nodes: '*'\n commands: ['show version']\n output: stdout" + }) + app.services.ai.build_playbook_chat = mock_chat + + with patch("rich.prompt.Prompt.ask", side_effect=["create task", "run"]): + with patch("connpy.cli.run_handler.RunHandler.yaml_run") as mock_yaml_run: + app.start(["run", "--generate-ai", str(dest_yaml)]) + + mock_chat.assert_called_once_with("create task", chat_history=[], chunk_callback=ANY) + assert os.path.exists(dest_yaml) + with open(dest_yaml) as f: + content = f.read() + assert "tasks:" in content + + mock_yaml_run.assert_called_once() + args = mock_yaml_run.call_args[0][0] + assert args.data == [str(dest_yaml)] diff --git a/connpy/tests/test_run_ai.py b/connpy/tests/test_run_ai.py new file mode 100644 index 0000000..20736ca --- /dev/null +++ b/connpy/tests/test_run_ai.py @@ -0,0 +1,296 @@ +import pytest +import json +from unittest.mock import patch, MagicMock +from connpy.ai import PlaybookBuilderAgent +from connpy.services.ai_service import AIService + +# ========================================================================= +# PlaybookBuilderAgent validation tests +# ========================================================================= + +def test_validate_playbook_valid(ai_config): + """Verifies that a valid canonical tasks[] playbook passes validation.""" + agent = PlaybookBuilderAgent(ai_config) + + valid_yaml = """ + tasks: + - name: "Apply standard config" + action: "run" + nodes: "router1" + commands: + - "conf t" + - "end" + output: "stdout" + - name: "Verify connectivity" + action: "test" + nodes: "router1" + commands: + - "ping 10.0.0.1" + expected: "!" + output: "stdout" + """ + + res = agent.validate_playbook(valid_yaml) + assert res["valid"] is True + assert "valid" in res["message"].lower() + +def test_validate_playbook_invalid_yaml(ai_config): + """Verifies that syntax errors in YAML are caught and reported.""" + agent = PlaybookBuilderAgent(ai_config) + + invalid_yaml = """ + tasks: + - name: "Broken task" + action: "run + nodes: "router1" + """ + + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "syntax error" in res["error"].lower() + +def test_validate_playbook_missing_tasks_key(ai_config): + """Verifies that a playbook without tasks root key is invalid.""" + agent = PlaybookBuilderAgent(ai_config) + + invalid_yaml = """ + not_tasks: + - name: "Apply standard config" + action: "run" + nodes: "router1" + commands: + - "conf t" + output: "stdout" + """ + + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "missing mandatory root 'tasks' key" in res["error"].lower() + +def test_validate_playbook_missing_mandatory_fields(ai_config): + """Verifies that missing name, action, nodes, commands, or output triggers a validation failure.""" + agent = PlaybookBuilderAgent(ai_config) + + # Missing nodes + invalid_yaml = """ + tasks: + - name: "Apply standard config" + action: "run" + commands: + - "conf t" + output: "stdout" + """ + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "missing mandatory fields" in res["error"].lower() + assert "nodes" in res["error"] + +def test_validate_playbook_invalid_action(ai_config): + """Verifies that an unsupported action type is caught.""" + agent = PlaybookBuilderAgent(ai_config) + + invalid_yaml = """ + tasks: + - name: "Apply standard config" + action: "delete_everything" + nodes: "router1" + commands: + - "conf t" + output: "stdout" + """ + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "invalid action" in res["error"].lower() + +def test_validate_playbook_missing_expected_in_test(ai_config): + """Verifies that action 'test' requires the expected field.""" + agent = PlaybookBuilderAgent(ai_config) + + invalid_yaml = """ + tasks: + - name: "Apply standard config" + action: "test" + nodes: "router1" + commands: + - "ping 10.0.0.1" + output: "stdout" + """ + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "missing the mandatory 'expected' key" in res["error"].lower() + +def test_validate_playbook_invalid_nodes_type(ai_config): + """Verifies that nodes of invalid type (e.g. integer) is caught.""" + agent = PlaybookBuilderAgent(ai_config) + + invalid_yaml = """ + tasks: + - name: "Apply config" + action: "run" + nodes: 12345 + commands: + - "conf t" + output: "stdout" + """ + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "nodes' must be a string (regex) or a list of strings (regexes)" in res["error"] + +def test_validate_playbook_invalid_nodes_list_item(ai_config): + """Verifies that nodes list containing non-string items is caught.""" + agent = PlaybookBuilderAgent(ai_config) + + invalid_yaml = """ + tasks: + - name: "Apply config" + action: "run" + nodes: + - "router1" + - 9999 + commands: + - "conf t" + output: "stdout" + """ + res = agent.validate_playbook(invalid_yaml) + assert res["valid"] is False + assert "list contains a non-string value" in res["error"] + + +# ========================================================================= +# AIService new methods delegation tests +# ========================================================================= + +def test_build_playbook_chat_delegation(ai_config): + """Verifies that build_playbook_chat instantiates PlaybookBuilderAgent and delegates ask.""" + service = AIService(ai_config) + + with patch("connpy.ai.PlaybookBuilderAgent") as MockAgentClass: + mock_agent = MockAgentClass.return_value + mock_agent.ask.return_value = {"response": "Mock response", "chat_history": []} + + history = [{"role": "user", "content": "build playbook"}] + res = service.build_playbook_chat("help me", chat_history=history) + + MockAgentClass.assert_called_once_with(ai_config) + mock_agent.ask.assert_called_once_with("help me", chat_history=history, status=None, chunk_callback=None) + assert res["response"] == "Mock response" + +def test_analyze_execution_results_delegation(ai_config): + """Verifies that analyze_execution_results formats prompt with @architect and delegates to self.ask.""" + service = AIService(ai_config) + service.ask = MagicMock() + + results = {"router1": {"output": "success", "status": 0}} + service.analyze_execution_results(results, query="diagnose border") + + service.ask.assert_called_once() + args, kwargs = service.ask.call_args + prompt = args[0] + + assert prompt.startswith("@architect:") + assert "diagnose border" in prompt + assert "Results Data:" in prompt + assert "router1" in prompt + assert kwargs.get("one_shot") is True + +def test_predict_execution_results_delegation(ai_config): + """Verifies that predict_execution_results formats prompt with @engineer and delegates to self.ask.""" + service = AIService(ai_config) + service.ask = MagicMock() + + nodes = ["router1", "router2"] + commands = ["conf t", "interface lo0"] + service.predict_execution_results(nodes, commands) + + service.ask.assert_called_once() + args, kwargs = service.ask.call_args + prompt = args[0] + + assert prompt.startswith("@engineer:") + assert "Preflight Simulation Agent" in prompt + assert "router1, router2" in prompt + assert "conf t" in prompt + assert "interface lo0" in prompt + + +# ========================================================================= +# gRPC Integration Tests for AIService +# ========================================================================= + +import grpc +from concurrent import futures +from connpy.grpc_layer import server, connpy_pb2, connpy_pb2_grpc, stubs + +class TestGRPCAIIntegration: + @pytest.fixture + def grpc_server(self, populated_config): + """Starts a local gRPC server for IA integration testing.""" + srv = grpc.server(futures.ThreadPoolExecutor(max_workers=5)) + connpy_pb2_grpc.add_AIServiceServicer_to_server(server.ServerServicer(populated_config).ai if hasattr(server, 'ServerServicer') else server.AIServicer(populated_config), srv) + port = srv.add_insecure_port('127.0.0.1:0') + srv.start() + yield f"127.0.0.1:{port}" + srv.stop(0) + + @pytest.fixture + def channel(self, grpc_server): + with grpc.insecure_channel(grpc_server) as channel: + yield channel + + @pytest.fixture + def ai_stub(self, channel): + return stubs.AIStub(channel, "localhost") + + def test_build_playbook_chat_grpc(self, ai_stub, populated_config): + """Verifies that build_playbook_chat gRPC stream functions correctly.""" + # Mock PlaybookBuilderAgent.ask to simulate agent response stream + def mock_ask(user_input, chat_history=None, status=None, debug=False, chunk_callback=None): + if chunk_callback: + chunk_callback("Generated Tasks:\n- name: config") + return {"response": "Done", "playbook_yaml": "tasks:\n- name: config"} + + with patch("connpy.ai.PlaybookBuilderAgent.ask", side_effect=mock_ask): + chunks = [] + def callback(chunk): + chunks.append(chunk) + + res = ai_stub.build_playbook_chat("make playbook", chunk_callback=callback) + assert "tasks:" in res["playbook_yaml"] + assert len(chunks) > 0 + assert "Generated Tasks:" in chunks[0] + + def test_analyze_execution_results_grpc(self, ai_stub, populated_config): + """Verifies that analyze_execution_results gRPC stream functions correctly.""" + # Mock AIService.ask to simulate response stream + def mock_ask(prompt, status=None, debug=False, chunk_callback=None, **kwargs): + if chunk_callback: + chunk_callback("Results are optimal.") + return {"response": "Done"} + + with patch.object(AIService, "ask", side_effect=mock_ask): + chunks = [] + def callback(chunk): + chunks.append(chunk) + + res = ai_stub.analyze_execution_results({"r1": "ok"}, query="test query", chunk_callback=callback) + assert res is not None + assert len(chunks) > 0 + assert "optimal" in chunks[0] + + def test_predict_execution_results_grpc(self, ai_stub, populated_config): + """Verifies that predict_execution_results gRPC stream functions correctly.""" + # Mock AIService.ask to simulate response stream + def mock_ask(prompt, status=None, debug=False, chunk_callback=None, **kwargs): + if chunk_callback: + chunk_callback("Commands are safe.") + return {"response": "Done"} + + with patch.object(AIService, "ask", side_effect=mock_ask): + chunks = [] + def callback(chunk): + chunks.append(chunk) + + res = ai_stub.predict_execution_results(["r1"], ["show version"], chunk_callback=callback) + assert res is not None + assert len(chunks) > 0 + assert "safe" in chunks[0] diff --git a/docs/connpy/cli/ai_handler.html b/docs/connpy/cli/ai_handler.html index b50170c..740247e 100644 --- a/docs/connpy/cli/ai_handler.html +++ b/docs/connpy/cli/ai_handler.html @@ -140,7 +140,7 @@ el.replaceWith(d); def single_question(self, args, session_id): query = " ".join(args.ask) - with console.status("[ai_status]Agent is thinking and analyzing...") as status: + with console.status("[ai_status]Agent is thinking and analyzing...[/ai_status]") as status: result = self.app.myai.ask(query, status=status, debug=args.debug, session_id=session_id, trust=args.trust, **self.ai_overrides) responder = result.get("responder", "engineer") @@ -177,7 +177,7 @@ el.replaceWith(d); if not user_query.strip(): continue if user_query.lower() in ['exit', 'quit', 'bye', 'cancel']: break - with console.status("[ai_status]Agent is thinking...") as status: + with console.status("[ai_status]Agent is thinking...[/ai_status]") as status: result = self.app.myai.ask(user_query, chat_history=history, status=status, debug=args.debug, trust=args.trust, session_id=session_id, **self.ai_overrides) new_history = result.get("chat_history") @@ -583,7 +583,7 @@ el.replaceWith(d); if not user_query.strip(): continue if user_query.lower() in ['exit', 'quit', 'bye', 'cancel']: break - with console.status("[ai_status]Agent is thinking...") as status: + with console.status("[ai_status]Agent is thinking...[/ai_status]") as status: result = self.app.myai.ask(user_query, chat_history=history, status=status, debug=args.debug, trust=args.trust, session_id=session_id, **self.ai_overrides) new_history = result.get("chat_history") @@ -618,7 +618,7 @@ el.replaceWith(d); 
def single_question(self, args, session_id):
     query = " ".join(args.ask)
-    with console.status("[ai_status]Agent is thinking and analyzing...") as status:
+    with console.status("[ai_status]Agent is thinking and analyzing...[/ai_status]") as status:
         result = self.app.myai.ask(query, status=status, debug=args.debug, session_id=session_id, trust=args.trust, **self.ai_overrides)
     
     responder = result.get("responder", "engineer")
diff --git a/docs/connpy/cli/run_handler.html b/docs/connpy/cli/run_handler.html
index 11759c3..9492172 100644
--- a/docs/connpy/cli/run_handler.html
+++ b/docs/connpy/cli/run_handler.html
@@ -63,7 +63,12 @@ el.replaceWith(d);
     def dispatch(self, args):
         if len(args.data) > 1:
             args.action = "noderun"
-        actions = {"noderun": self.node_run, "generate": self.yaml_generate, "run": self.yaml_run}
+        actions = {
+            "noderun": self.node_run,
+            "generate": self.yaml_generate,
+            "generate_ai": self.ai_generate,
+            "run": self.yaml_run
+        }
         return actions.get(args.action)(args)
 
     def node_run(self, args):
@@ -81,6 +86,41 @@ el.replaceWith(d);
             
         commands = [" ".join(args.data[1:])]
 
+        # Check for Preflight AI simulation
+        if getattr(args, "preflight_ai", False):
+            matched_node_names = [n.get("name") if isinstance(n, dict) else n for n in matched_nodes]
+            
+            renderer = printer.BlockMarkdownRenderer()
+            first_chunk = True
+            status_context = printer.console.status("[ai_status]Simulating execution...[/ai_status]")
+            
+            def callback(chunk):
+                nonlocal first_chunk
+                if first_chunk:
+                    try: status_context.stop()
+                    except: pass
+                    printer.console.print(Rule(title="[engineer][bold]Preflight AI Simulation[/bold][/engineer]", style="engineer"))
+                    first_chunk = False
+                renderer.feed(chunk)
+            
+            try:
+                status_context.start()
+                self.app.services.ai.predict_execution_results(
+                    matched_node_names,
+                    commands,
+                    chunk_callback=callback
+                )
+                if first_chunk:
+                    try: status_context.stop()
+                    except: pass
+                    printer.console.print(Rule(title="[engineer][bold]Preflight AI Simulation[/bold][/engineer]", style="engineer"))
+                renderer.flush()
+                printer.console.print(Rule(style="engineer"))
+            except Exception as e:
+                printer.error(f"Preflight AI simulation failed: {e}")
+                sys.exit(1)
+            sys.exit(0)
+
         try:
             header_printed = False
 
@@ -118,6 +158,40 @@ el.replaceWith(d);
                 )
                 printer.run_summary(results)
 
+            # Analyze execution results if requested
+            if getattr(args, "analyze", None) is not None:
+                printer.console.print()
+                
+                renderer = printer.BlockMarkdownRenderer()
+                first_chunk = True
+                status_context = printer.console.status("[ai_status]Analyzing execution results...[/ai_status]")
+                
+                def callback(chunk):
+                    nonlocal first_chunk
+                    if first_chunk:
+                        try: status_context.stop()
+                        except: pass
+                        printer.console.print(Rule(title="[architect][bold]Network Architect AI Analysis[/bold][/architect]", style="architect"))
+                        first_chunk = False
+                    renderer.feed(chunk)
+                
+                query = args.analyze if args.analyze else " ".join(args.data[1:])
+                try:
+                    status_context.start()
+                    self.app.services.ai.analyze_execution_results(
+                        results,
+                        query=query,
+                        chunk_callback=callback
+                    )
+                    if first_chunk:
+                        try: status_context.stop()
+                        except: pass
+                        printer.console.print(Rule(title="[architect][bold]Network Architect AI Analysis[/bold][/architect]", style="architect"))
+                    renderer.flush()
+                    printer.console.print(Rule(style="architect"))
+                except Exception as e:
+                    printer.error(f"AI Analysis failed: {e}")
+
         except ConnpyError as e:
             printer.error(str(e))
             sys.exit(1)
@@ -138,8 +212,105 @@ el.replaceWith(d);
             with open(path, "r") as f:
                 playbook = yaml.load(f, Loader=yaml.FullLoader)
 
+            # Check preflight first before any task runs
+            if getattr(args, "preflight_ai", False):
+                preflight_failed = False
+                for task in playbook.get("tasks", []):
+                    name = task.get("name", "Task")
+                    nodelist = task.get("nodes", [])
+                    commands = task.get("commands", [])
+                    
+                    # Resolve nodes to names
+                    try:
+                        if isinstance(nodelist, str):
+                            resolved_nodes = self.app.services.nodes.list_nodes(nodelist)
+                        elif isinstance(nodelist, list):
+                            resolved_nodes = []
+                            for item in nodelist:
+                                matches = self.app.services.nodes.list_nodes(item)
+                                for m in matches:
+                                    if m not in resolved_nodes:
+                                        resolved_nodes.append(m)
+                        else:
+                            resolved_nodes = []
+                    except Exception:
+                        resolved_nodes = []
+                    
+                    resolved_names = [n.get("name") if isinstance(n, dict) else n for n in resolved_nodes]
+                    printer.console.print(f"\n[bold]Task: {name}[/bold] (Preflight for {len(resolved_names)} nodes)")
+                    
+                    renderer = printer.BlockMarkdownRenderer()
+                    first_chunk = True
+                    status_context = printer.console.status("[ai_status]Simulating execution...[/ai_status]")
+                    
+                    def callback(chunk):
+                        nonlocal first_chunk
+                        if first_chunk:
+                            try: status_context.stop()
+                            except: pass
+                            printer.console.print(Rule(title=f"[engineer][bold]Preflight AI Simulation: {name}[/bold][/engineer]", style="engineer"))
+                            first_chunk = False
+                        renderer.feed(chunk)
+                    try:
+                        status_context.start()
+                        self.app.services.ai.predict_execution_results(
+                            resolved_names,
+                            commands,
+                            chunk_callback=callback
+                        )
+                        if first_chunk:
+                            try: status_context.stop()
+                            except: pass
+                            printer.console.print(Rule(title=f"[engineer][bold]Preflight AI Simulation: {name}[/bold][/engineer]", style="engineer"))
+                        renderer.flush()
+                        printer.console.print(Rule(style="engineer"))
+                    except Exception as e:
+                        printer.error(f"Preflight AI simulation failed for task {name}: {e}")
+                        preflight_failed = True
+                if preflight_failed:
+                    sys.exit(1)
+                sys.exit(0)
+
+            # Standard run
+            results_all = {}
             for task in playbook.get("tasks", []):
-                self.cli_run(task)
+                task_res = self.cli_run(task)
+                if task_res:
+                    results_all.update(task_res)
+
+            # If analyze is enabled, run analysis on accumulated results
+            if getattr(args, "analyze", None) is not None:
+                printer.console.print()
+                
+                renderer = printer.BlockMarkdownRenderer()
+                first_chunk = True
+                status_context = printer.console.status("[ai_status]Analyzing playbook execution results...[/ai_status]")
+                
+                def callback(chunk):
+                    nonlocal first_chunk
+                    if first_chunk:
+                        try: status_context.stop()
+                        except: pass
+                        printer.console.print(Rule(title="[architect][bold]Network Architect AI Playbook Analysis[/bold][/architect]", style="architect"))
+                        first_chunk = False
+                    renderer.feed(chunk)
+                
+                query = args.analyze if args.analyze else f"Playbook: {path}"
+                try:
+                    status_context.start()
+                    self.app.services.ai.analyze_execution_results(
+                        results_all,
+                        query=query,
+                        chunk_callback=callback
+                    )
+                    if first_chunk:
+                        try: status_context.stop()
+                        except: pass
+                        printer.console.print(Rule(title="[architect][bold]Network Architect AI Playbook Analysis[/bold][/architect]", style="architect"))
+                    renderer.flush()
+                    printer.console.print(Rule(style="architect"))
+                except Exception as e:
+                    printer.error(f"AI Analysis failed: {e}")
 
         except Exception as e:
             printer.error(f"Failed to run playbook {path}: {e}")
@@ -184,6 +355,7 @@ el.replaceWith(d);
 
         nodelist = resolved_nodes
 
+        results = {}
         try:
             header_printed = False
             if action == "run":
@@ -243,13 +415,244 @@ el.replaceWith(d);
                 )
                 # ALWAYS show the aggregate summary at the end
                 printer.test_summary(results)
+                
+            return results
 
         except ConnpyError as e:
-            printer.error(str(e))
+ printer.error(str(e)) + return {} + + def ai_generate(self, args): + from rich.prompt import Prompt + from rich.rule import Rule + from rich.panel import Panel + from rich.syntax import Syntax + + dest_file = args.data[0] + if os.path.exists(dest_file): + printer.error(f"File '{dest_file}' already exists.") + sys.exit(14) + + chat_history = [] + + # Consistent layout opening matching global AI (engineer style) + from rich.markdown import Markdown + printer.console.print(Rule(style="engineer")) + printer.console.print(Markdown("**Playbook Builder AI**: Welcome! Describe the automation workflow you want to design.\nType **exit** to quit.\n")) + printer.console.print(Rule(style="engineer")) + + while True: + try: + user_prompt = Prompt.ask("[user_prompt]User[/user_prompt]") + except (KeyboardInterrupt, EOFError): + printer.console.print() + printer.warning("Operation cancelled by user.") + break + + if user_prompt.strip().lower() in ["exit", "quit"]: + printer.info("Exiting AI Assistant.") + break + + if not user_prompt.strip(): + continue + + printer.console.print() + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Agent is thinking...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: + status_context.stop() + except: + pass + printer.console.print(Rule(title="[engineer][bold]Playbook Builder AI[/bold][/engineer]", style="engineer")) + first_chunk = False + renderer.feed(chunk) + + try: + status_context.start() + res = self.app.services.ai.build_playbook_chat( + user_prompt, + chat_history=chat_history, + chunk_callback=callback + ) + if first_chunk: + try: + status_context.stop() + except: + pass + renderer.flush() + if not first_chunk: + printer.console.print(Rule(style="engineer")) + + # Update history + if res and "chat_history" in res: + chat_history = res["chat_history"] + + # Check if the agent returned a validated playbook YAML + if res and "playbook_yaml" in res and res["playbook_yaml"]: + yaml_content = res["playbook_yaml"] + printer.console.print() + printer.success("Playbook YAML successfully generated and validated.") + + # Show the YAML inside a beautiful panel matching AI style (with engineer borders) + syntax = Syntax(yaml_content, "yaml", theme="ansi_dark", word_wrap=True, background_color="default") + panel = Panel(syntax, title="[engineer][bold]Resulting Playbook[/bold][/engineer]", border_style="engineer", expand=False) + printer.console.print(panel) + + # Ask if the user wants to save it + try: + save_confirm = Prompt.ask( + f"\nDo you want to save this playbook to '{dest_file}'?", + choices=["y", "n", "run"], + default="y" + ) + except (KeyboardInterrupt, EOFError): + printer.console.print() + printer.warning("Saving skipped.") + break + + choice = save_confirm.strip().lower() + if choice in ["y", "yes", "run"]: + with open(dest_file, "w") as f: + f.write(yaml_content) + printer.success(f"Playbook saved successfully to '{dest_file}'") + if choice == "run": + printer.console.print() + printer.info("Executing the saved playbook...") + self.yaml_run(args) + break + else: + printer.warning("Playbook not saved. You can continue describing changes or exit.") + except Exception as e: + printer.error(f"Error in AI chat: {e}")

Methods

+
+def ai_generate(self, args) +
+
+
+ +Expand source code + +
def ai_generate(self, args):
+    from rich.prompt import Prompt
+    from rich.rule import Rule
+    from rich.panel import Panel
+    from rich.syntax import Syntax
+
+    dest_file = args.data[0]
+    if os.path.exists(dest_file):
+        printer.error(f"File '{dest_file}' already exists.")
+        sys.exit(14)
+
+    chat_history = []
+    
+    # Consistent layout opening matching global AI (engineer style)
+    from rich.markdown import Markdown
+    printer.console.print(Rule(style="engineer"))
+    printer.console.print(Markdown("**Playbook Builder AI**: Welcome! Describe the automation workflow you want to design.\nType **exit** to quit.\n"))
+    printer.console.print(Rule(style="engineer"))
+    
+    while True:
+        try:
+            user_prompt = Prompt.ask("[user_prompt]User[/user_prompt]")
+        except (KeyboardInterrupt, EOFError):
+            printer.console.print()
+            printer.warning("Operation cancelled by user.")
+            break
+            
+        if user_prompt.strip().lower() in ["exit", "quit"]:
+            printer.info("Exiting AI Assistant.")
+            break
+            
+        if not user_prompt.strip():
+            continue
+            
+        printer.console.print()
+        
+        renderer = printer.BlockMarkdownRenderer()
+        first_chunk = True
+        status_context = printer.console.status("[ai_status]Agent is thinking...[/ai_status]")
+        
+        def callback(chunk):
+            nonlocal first_chunk
+            if first_chunk:
+                try:
+                    status_context.stop()
+                except:
+                    pass
+                printer.console.print(Rule(title="[engineer][bold]Playbook Builder AI[/bold][/engineer]", style="engineer"))
+                first_chunk = False
+            renderer.feed(chunk)
+            
+        try:
+            status_context.start()
+            res = self.app.services.ai.build_playbook_chat(
+                user_prompt,
+                chat_history=chat_history,
+                chunk_callback=callback
+            )
+            if first_chunk:
+                try:
+                    status_context.stop()
+                except:
+                    pass
+            renderer.flush()
+            if not first_chunk:
+                printer.console.print(Rule(style="engineer"))
+            
+            # Update history
+            if res and "chat_history" in res:
+                chat_history = res["chat_history"]
+            
+            # Check if the agent returned a validated playbook YAML
+            if res and "playbook_yaml" in res and res["playbook_yaml"]:
+                yaml_content = res["playbook_yaml"]
+                printer.console.print()
+                printer.success("Playbook YAML successfully generated and validated.")
+                
+                # Show the YAML inside a beautiful panel matching AI style (with engineer borders)
+                syntax = Syntax(yaml_content, "yaml", theme="ansi_dark", word_wrap=True, background_color="default")
+                panel = Panel(syntax, title="[engineer][bold]Resulting Playbook[/bold][/engineer]", border_style="engineer", expand=False)
+                printer.console.print(panel)
+                
+                # Ask if the user wants to save it
+                try:
+                    save_confirm = Prompt.ask(
+                        f"\nDo you want to save this playbook to '{dest_file}'?",
+                        choices=["y", "n", "run"],
+                        default="y"
+                    )
+                except (KeyboardInterrupt, EOFError):
+                    printer.console.print()
+                    printer.warning("Saving skipped.")
+                    break
+                    
+                choice = save_confirm.strip().lower()
+                if choice in ["y", "yes", "run"]:
+                    with open(dest_file, "w") as f:
+                        f.write(yaml_content)
+                    printer.success(f"Playbook saved successfully to '{dest_file}'")
+                    if choice == "run":
+                        printer.console.print()
+                        printer.info("Executing the saved playbook...")
+                        self.yaml_run(args)
+                    break
+                else:
+                    printer.warning("Playbook not saved. You can continue describing changes or exit.")
+        except Exception as e:
+            printer.error(f"Error in AI chat: {e}")
+
+
+
def cli_run(self, script)
@@ -297,6 +700,7 @@ el.replaceWith(d); nodelist = resolved_nodes + results = {} try: header_printed = False if action == "run": @@ -356,9 +760,12 @@ el.replaceWith(d); ) # ALWAYS show the aggregate summary at the end printer.test_summary(results) + + return results except ConnpyError as e: - printer.error(str(e)) + printer.error(str(e)) + return {}
@@ -373,7 +780,12 @@ el.replaceWith(d); 
def dispatch(self, args):
     if len(args.data) > 1:
         args.action = "noderun"
-    actions = {"noderun": self.node_run, "generate": self.yaml_generate, "run": self.yaml_run}
+    actions = {
+        "noderun": self.node_run,
+        "generate": self.yaml_generate,
+        "generate_ai": self.ai_generate,
+        "run": self.yaml_run
+    }
     return actions.get(args.action)(args)
@@ -401,6 +813,41 @@ el.replaceWith(d); commands = [" ".join(args.data[1:])] + # Check for Preflight AI simulation + if getattr(args, "preflight_ai", False): + matched_node_names = [n.get("name") if isinstance(n, dict) else n for n in matched_nodes] + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Simulating execution...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[engineer][bold]Preflight AI Simulation[/bold][/engineer]", style="engineer")) + first_chunk = False + renderer.feed(chunk) + + try: + status_context.start() + self.app.services.ai.predict_execution_results( + matched_node_names, + commands, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[engineer][bold]Preflight AI Simulation[/bold][/engineer]", style="engineer")) + renderer.flush() + printer.console.print(Rule(style="engineer")) + except Exception as e: + printer.error(f"Preflight AI simulation failed: {e}") + sys.exit(1) + sys.exit(0) + try: header_printed = False @@ -438,6 +885,40 @@ el.replaceWith(d); ) printer.run_summary(results) + # Analyze execution results if requested + if getattr(args, "analyze", None) is not None: + printer.console.print() + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Analyzing execution results...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Analysis[/bold][/architect]", style="architect")) + first_chunk = False + renderer.feed(chunk) + + query = args.analyze if args.analyze else " ".join(args.data[1:]) + try: + status_context.start() + self.app.services.ai.analyze_execution_results( + results, + query=query, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Analysis[/bold][/architect]", style="architect")) + renderer.flush() + printer.console.print(Rule(style="architect")) + except Exception as e: + printer.error(f"AI Analysis failed: {e}") + except ConnpyError as e: printer.error(str(e)) sys.exit(1) @@ -478,8 +959,105 @@ el.replaceWith(d); with open(path, "r") as f: playbook = yaml.load(f, Loader=yaml.FullLoader) + # Check preflight first before any task runs + if getattr(args, "preflight_ai", False): + preflight_failed = False + for task in playbook.get("tasks", []): + name = task.get("name", "Task") + nodelist = task.get("nodes", []) + commands = task.get("commands", []) + + # Resolve nodes to names + try: + if isinstance(nodelist, str): + resolved_nodes = self.app.services.nodes.list_nodes(nodelist) + elif isinstance(nodelist, list): + resolved_nodes = [] + for item in nodelist: + matches = self.app.services.nodes.list_nodes(item) + for m in matches: + if m not in resolved_nodes: + resolved_nodes.append(m) + else: + resolved_nodes = [] + except Exception: + resolved_nodes = [] + + resolved_names = [n.get("name") if isinstance(n, dict) else n for n in resolved_nodes] + printer.console.print(f"\n[bold]Task: {name}[/bold] (Preflight for {len(resolved_names)} nodes)") + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Simulating execution...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title=f"[engineer][bold]Preflight AI Simulation: {name}[/bold][/engineer]", style="engineer")) + first_chunk = False + renderer.feed(chunk) + try: + status_context.start() + self.app.services.ai.predict_execution_results( + resolved_names, + commands, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title=f"[engineer][bold]Preflight AI Simulation: {name}[/bold][/engineer]", style="engineer")) + renderer.flush() + printer.console.print(Rule(style="engineer")) + except Exception as e: + printer.error(f"Preflight AI simulation failed for task {name}: {e}") + preflight_failed = True + if preflight_failed: + sys.exit(1) + sys.exit(0) + + # Standard run + results_all = {} for task in playbook.get("tasks", []): - self.cli_run(task) + task_res = self.cli_run(task) + if task_res: + results_all.update(task_res) + + # If analyze is enabled, run analysis on accumulated results + if getattr(args, "analyze", None) is not None: + printer.console.print() + + renderer = printer.BlockMarkdownRenderer() + first_chunk = True + status_context = printer.console.status("[ai_status]Analyzing playbook execution results...[/ai_status]") + + def callback(chunk): + nonlocal first_chunk + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Playbook Analysis[/bold][/architect]", style="architect")) + first_chunk = False + renderer.feed(chunk) + + query = args.analyze if args.analyze else f"Playbook: {path}" + try: + status_context.start() + self.app.services.ai.analyze_execution_results( + results_all, + query=query, + chunk_callback=callback + ) + if first_chunk: + try: status_context.stop() + except: pass + printer.console.print(Rule(title="[architect][bold]Network Architect AI Playbook Analysis[/bold][/architect]", style="architect")) + renderer.flush() + printer.console.print(Rule(style="architect")) + except Exception as e: + printer.error(f"AI Analysis failed: {e}") except Exception as e: printer.error(f"Failed to run playbook {path}: {e}") @@ -506,7 +1084,8 @@ el.replaceWith(d);