From e78f61d6d068904c5dd6df298a4a5015ffa2816c Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 18 Feb 2026 04:15:40 +0000 Subject: [PATCH] feat(python): add ScriptedTool bindings and LangChain integration - Add ScriptedTool Python class wrapping Rust ScriptedTool via PyO3 - add_tool() registers Python callbacks as bash builtins - Callbacks receive (params: dict, stdin: str | None) -> str - JSON<->Python conversion for typed flag parsing (int, bool, str) - execute/execute_sync run bash scripts with all registered tools - Introspection: system_prompt(), help(), description(), schemas - Add ScriptedToolLangChain wrapper and create_scripted_tool() factory - Add k8s_orchestrator.py example: 12 fake kubectl commands, 6 demos - get_nodes, get_namespaces, get_pods, get_deployments, get_services - describe_pod, get_logs, get_configmaps, get_secrets, get_events - scale_deployment, rollout_status - Optional LangChain ReAct agent integration (--langchain flag) - Add 26 Python tests covering construction, execution, pipelines, jq integration, error handling, stdin, env vars, loops, conditionals, async, introspection, and 12-tool orchestration - Enable scripted_tool feature in bashkit-python Cargo.toml - Update type stubs and __init__.py exports https://claude.ai/code/session_0117wzNc2aRc5chtBgfwgdCd --- crates/bashkit-python/Cargo.toml | 2 +- crates/bashkit-python/bashkit/__init__.py | 13 +- crates/bashkit-python/bashkit/_bashkit.pyi | 104 ++-- crates/bashkit-python/bashkit/langchain.py | 131 ++++- .../examples/k8s_orchestrator.py | 526 ++++++++++++++++++ crates/bashkit-python/src/lib.rs | 422 +++++++++++--- crates/bashkit-python/tests/test_bashkit.py | 318 ++++++++++- 7 files changed, 1370 insertions(+), 146 deletions(-) create mode 100644 crates/bashkit-python/examples/k8s_orchestrator.py diff --git a/crates/bashkit-python/Cargo.toml b/crates/bashkit-python/Cargo.toml index a49307d9..d17b4997 100644 --- a/crates/bashkit-python/Cargo.toml +++ b/crates/bashkit-python/Cargo.toml @@ -17,7 +17,7 @@ doc = false # Python extension, no Rust docs needed [dependencies] # Bashkit core -bashkit = { path = "../bashkit" } +bashkit = { path = "../bashkit", features = ["scripted_tool"] } # PyO3 for Python bindings pyo3 = { workspace = true } diff --git a/crates/bashkit-python/bashkit/__init__.py b/crates/bashkit-python/bashkit/__init__.py index 3301b60a..f5fa7c14 100644 --- a/crates/bashkit-python/bashkit/__init__.py +++ b/crates/bashkit-python/bashkit/__init__.py @@ -10,22 +10,26 @@ >>> print(result.stdout) Hello, World! +For scripted multi-tool orchestration: + >>> from bashkit import ScriptedTool + >>> tool = ScriptedTool("api") + >>> tool.add_tool("greet", "Greet user", callback=lambda p, s=None: f"hello {p.get('name', 'world')}") + >>> result = tool.execute_sync("greet --name Alice") + For LangChain integration: - >>> from bashkit.langchain import create_bash_tool - >>> tool = create_bash_tool() + >>> from bashkit.langchain import create_bash_tool, create_scripted_tool For Deep Agents integration: >>> from bashkit.deepagents import create_bash_middleware - >>> middleware = create_bash_middleware() For PydanticAI integration: >>> from bashkit.pydantic_ai import create_bash_tool - >>> tool = create_bash_tool() """ from bashkit._bashkit import ( BashTool, ExecResult, + ScriptedTool, create_langchain_tool_spec, ) @@ -33,5 +37,6 @@ __all__ = [ "BashTool", "ExecResult", + "ScriptedTool", "create_langchain_tool_spec", ] diff --git a/crates/bashkit-python/bashkit/_bashkit.pyi b/crates/bashkit-python/bashkit/_bashkit.pyi index ce9380e9..66972ba9 100644 --- a/crates/bashkit-python/bashkit/_bashkit.pyi +++ b/crates/bashkit-python/bashkit/_bashkit.pyi @@ -1,5 +1,7 @@ """Type stubs for bashkit_py native module.""" +from typing import Any, Callable + class ExecResult: """Result from executing bash commands.""" @@ -9,7 +11,7 @@ class ExecResult: error: str | None success: bool - def to_dict(self) -> dict[str, any]: ... + def to_dict(self) -> dict[str, Any]: ... class BashTool: """Sandboxed bash interpreter for AI agents. @@ -35,62 +37,60 @@ class BashTool: hostname: str | None = None, max_commands: int | None = None, max_loop_iterations: int | None = None, - ) -> None: - """Create a new BashTool instance. - - Args: - username: Custom username for sandbox (default: "user") - hostname: Custom hostname for sandbox (default: "sandbox") - max_commands: Maximum commands to execute (default: 10000) - max_loop_iterations: Maximum loop iterations (default: 100000) - """ - ... - - async def execute(self, commands: str) -> ExecResult: - """Execute bash commands asynchronously. - - Args: - commands: Bash commands to execute (like `bash -c "commands"`) - - Returns: - ExecResult with stdout, stderr, exit_code - """ - ... - - def execute_sync(self, commands: str) -> ExecResult: - """Execute bash commands synchronously (blocking). - - Note: Prefer `execute()` for async contexts. This method blocks. - - Args: - commands: Bash commands to execute - - Returns: - ExecResult with stdout, stderr, exit_code - """ - ... - - def description(self) -> str: - """Get the full description.""" - ... + ) -> None: ... + async def execute(self, commands: str) -> ExecResult: ... + def execute_sync(self, commands: str) -> ExecResult: ... + def description(self) -> str: ... + def help(self) -> str: ... + def system_prompt(self) -> str: ... + def input_schema(self) -> str: ... + def output_schema(self) -> str: ... - def help(self) -> str: - """Get LLM documentation.""" - ... +class ScriptedTool: + """Compose Python callbacks as bash builtins for multi-tool orchestration. - def system_prompt(self) -> str: - """Get system prompt for LLMs.""" - ... + Each registered tool becomes a bash builtin command. An LLM (or user) + writes a single bash script that pipes, loops, and branches across tools. - def input_schema(self) -> str: - """Get JSON schema for input validation.""" - ... + Example: + >>> tool = ScriptedTool("api") + >>> tool.add_tool("greet", "Greet user", + ... callback=lambda p, s=None: f"hello {p.get('name', 'world')}\\n", + ... schema={"type": "object", "properties": {"name": {"type": "string"}}}) + >>> result = tool.execute_sync("greet --name Alice") + >>> print(result.stdout.strip()) + hello Alice + """ - def output_schema(self) -> str: - """Get JSON schema for output.""" - ... + name: str + short_description: str + version: str -def create_langchain_tool_spec() -> dict[str, any]: + def __init__( + self, + name: str, + short_description: str | None = None, + max_commands: int | None = None, + max_loop_iterations: int | None = None, + ) -> None: ... + def add_tool( + self, + name: str, + description: str, + callback: Callable[[dict[str, Any], str | None], str], + schema: dict[str, Any] | None = None, + ) -> None: ... + def env(self, key: str, value: str) -> None: ... + async def execute(self, commands: str) -> ExecResult: ... + def execute_sync(self, commands: str) -> ExecResult: ... + def tool_count(self) -> int: ... + def description(self) -> str: ... + def help(self) -> str: ... + def system_prompt(self) -> str: ... + def input_schema(self) -> str: ... + def output_schema(self) -> str: ... + +def create_langchain_tool_spec() -> dict[str, Any]: """Create a LangChain-compatible tool specification. Returns: diff --git a/crates/bashkit-python/bashkit/langchain.py b/crates/bashkit-python/bashkit/langchain.py index 6fd861fe..9c399a6c 100644 --- a/crates/bashkit-python/bashkit/langchain.py +++ b/crates/bashkit-python/bashkit/langchain.py @@ -1,15 +1,22 @@ """ LangChain integration for Bashkit. -Provides a LangChain-compatible tool that wraps BashTool for use with -LangChain agents and chains. +Provides LangChain-compatible tools wrapping BashTool and ScriptedTool +for use with LangChain agents and chains. -Example: +Example (BashTool): >>> from bashkit.langchain import create_bash_tool - >>> from langchain.agents import create_agent - >>> >>> tool = create_bash_tool() - >>> agent = create_agent(model="claude-sonnet-4-20250514", tools=[tool]) + >>> result = tool.invoke({"commands": "echo hello"}) + +Example (ScriptedTool): + >>> from bashkit import ScriptedTool + >>> from bashkit.langchain import create_scripted_tool + >>> + >>> st = ScriptedTool("api") + >>> st.add_tool("greet", "Greet user", callback=lambda p, s=None: f"hello {p.get('name')}\\n") + >>> tool = create_scripted_tool(st) + >>> result = tool.invoke({"commands": "greet --name Alice"}) """ from __future__ import annotations @@ -32,6 +39,7 @@ def PrivateAttr(*args, **kwargs): from bashkit import BashTool as NativeBashTool +from bashkit import ScriptedTool as NativeScriptedTool class BashToolInput(BaseModel): @@ -51,12 +59,11 @@ class BashkitTool(BaseTool): >>> print(result) # Hello! """ - name: str = "" # Set in __init__ from bashkit - description: str = "" # Set in __init__ from bashkit + name: str = "" + description: str = "" args_schema: type[BaseModel] = BashToolInput handle_tool_error: bool = True - # Internal state - use PrivateAttr for pydantic v2 compatibility _bash_tool: NativeBashTool = PrivateAttr() def __init__( @@ -67,21 +74,12 @@ def __init__( max_loop_iterations: int | None = None, **kwargs, ): - """Initialize BashkitTool. - - Args: - username: Custom username for sandbox - hostname: Custom hostname for sandbox - max_commands: Max commands to execute - max_loop_iterations: Max loop iterations - """ bash_tool = NativeBashTool( username=username, hostname=hostname, max_commands=max_commands, max_loop_iterations=max_loop_iterations, ) - # Use name and description from bashkit lib kwargs["name"] = bash_tool.name kwargs["description"] = bash_tool.description() super().__init__(**kwargs) @@ -94,7 +92,6 @@ def _run(self, commands: str) -> str: if result.error: raise ToolException(f"Execution error: {result.error}") - # Return combined output for the agent output = result.stdout if result.stderr: output += f"\nSTDERR: {result.stderr}" @@ -110,7 +107,64 @@ async def _arun(self, commands: str) -> str: if result.error: raise ToolException(f"Execution error: {result.error}") - # Return combined output for the agent + output = result.stdout + if result.stderr: + output += f"\nSTDERR: {result.stderr}" + if result.exit_code != 0: + output += f"\n[Exit code: {result.exit_code}]" + + return output + + class ScriptedToolLangChain(BaseTool): + """LangChain tool wrapper for Bashkit ScriptedTool. + + Wraps a pre-configured ScriptedTool (with registered Python callbacks) + as a LangChain tool. The LLM sends bash scripts that orchestrate all + registered sub-tools in one call. + + Example: + >>> from bashkit import ScriptedTool + >>> st = ScriptedTool("k8s") + >>> st.add_tool("get_pods", "List pods", callback=my_callback) + >>> tool = ScriptedToolLangChain(st) + >>> result = tool.invoke({"commands": "get_pods --namespace default | jq '.items | length'"}) + """ + + name: str = "" + description: str = "" + args_schema: type[BaseModel] = BashToolInput + handle_tool_error: bool = True + + _scripted_tool: NativeScriptedTool = PrivateAttr() + + def __init__(self, scripted_tool: NativeScriptedTool, **kwargs): + kwargs["name"] = scripted_tool.name + kwargs["description"] = scripted_tool.system_prompt() + super().__init__(**kwargs) + object.__setattr__(self, "_scripted_tool", scripted_tool) + + def _run(self, commands: str) -> str: + """Execute scripted tool commands synchronously.""" + result = self._scripted_tool.execute_sync(commands) + + if result.error: + raise ToolException(f"Execution error: {result.error}") + + output = result.stdout + if result.stderr: + output += f"\nSTDERR: {result.stderr}" + if result.exit_code != 0: + output += f"\n[Exit code: {result.exit_code}]" + + return output + + async def _arun(self, commands: str) -> str: + """Execute scripted tool commands asynchronously.""" + result = await self._scripted_tool.execute(commands) + + if result.error: + raise ToolException(f"Execution error: {result.error}") + output = result.stdout if result.stderr: output += f"\nSTDERR: {result.stderr}" @@ -158,4 +212,39 @@ def create_bash_tool( ) -__all__ = ["BashkitTool", "BashToolInput", "create_bash_tool"] +def create_scripted_tool(scripted_tool: NativeScriptedTool) -> ScriptedToolLangChain: + """Create a LangChain-compatible tool from a configured ScriptedTool. + + Args: + scripted_tool: A ScriptedTool with registered tool callbacks + + Returns: + ScriptedToolLangChain instance for use with LangChain agents + + Raises: + ImportError: If langchain-core is not installed + + Example: + >>> from bashkit import ScriptedTool + >>> from bashkit.langchain import create_scripted_tool + >>> + >>> st = ScriptedTool("api") + >>> st.add_tool("get_data", "Fetch data", callback=my_fn) + >>> tool = create_scripted_tool(st) + >>> # Use with: create_react_agent(model, [tool]) + """ + if not LANGCHAIN_AVAILABLE: + raise ImportError( + "langchain-core is required for LangChain integration. Install with: pip install 'bashkit[langchain]'" + ) + + return ScriptedToolLangChain(scripted_tool) + + +__all__ = [ + "BashkitTool", + "BashToolInput", + "ScriptedToolLangChain", + "create_bash_tool", + "create_scripted_tool", +] diff --git a/crates/bashkit-python/examples/k8s_orchestrator.py b/crates/bashkit-python/examples/k8s_orchestrator.py new file mode 100644 index 00000000..a7298c87 --- /dev/null +++ b/crates/bashkit-python/examples/k8s_orchestrator.py @@ -0,0 +1,526 @@ +#!/usr/bin/env python3 +"""Kubernetes API orchestrator using Bashkit ScriptedTool. + +Demonstrates composing 12 fake k8s API tools into a single ScriptedTool that +an LLM agent can call with bash scripts. Each tool becomes a bash builtin; +the agent writes one script to orchestrate them all. + +Run directly: + cd crates/bashkit-python && maturin develop && python examples/k8s_orchestrator.py + +With LangChain (optional): + pip install 'bashkit[langchain]' + ANTHROPIC_API_KEY=... python examples/k8s_orchestrator.py --langchain +""" + +from __future__ import annotations + +import json +import sys + +from bashkit import ScriptedTool + +# ============================================================================= +# Fake k8s data +# ============================================================================= + +NODES = [ + {"name": "node-1", "status": "Ready", "cpu": "4", "memory": "16Gi", "pods": 23}, + {"name": "node-2", "status": "Ready", "cpu": "8", "memory": "32Gi", "pods": 41}, + {"name": "node-3", "status": "NotReady", "cpu": "4", "memory": "16Gi", "pods": 0}, +] + +NAMESPACES = [ + {"name": "default", "status": "Active"}, + {"name": "kube-system", "status": "Active"}, + {"name": "monitoring", "status": "Active"}, + {"name": "production", "status": "Active"}, +] + +PODS = { + "default": [ + {"name": "web-abc12", "status": "Running", "restarts": 0, "node": "node-1", "image": "nginx:1.25"}, + {"name": "api-def34", "status": "Running", "restarts": 2, "node": "node-2", "image": "api:v2.1"}, + { + "name": "worker-ghi56", + "status": "CrashLoopBackOff", + "restarts": 15, + "node": "node-2", + "image": "worker:v1.0", + }, + ], + "kube-system": [ + {"name": "coredns-aaa11", "status": "Running", "restarts": 0, "node": "node-1", "image": "coredns:1.11"}, + {"name": "etcd-bbb22", "status": "Running", "restarts": 0, "node": "node-1", "image": "etcd:3.5"}, + ], + "monitoring": [ + {"name": "prometheus-ccc33", "status": "Running", "restarts": 0, "node": "node-2", "image": "prom:2.48"}, + {"name": "grafana-ddd44", "status": "Running", "restarts": 1, "node": "node-2", "image": "grafana:10.2"}, + ], + "production": [ + {"name": "app-eee55", "status": "Running", "restarts": 0, "node": "node-1", "image": "app:v3.2"}, + {"name": "app-fff66", "status": "Running", "restarts": 0, "node": "node-2", "image": "app:v3.2"}, + {"name": "db-ggg77", "status": "Pending", "restarts": 0, "node": "", "image": "postgres:16"}, + ], +} + +DEPLOYMENTS = { + "default": [ + {"name": "web", "replicas": 1, "available": 1, "image": "nginx:1.25"}, + {"name": "api", "replicas": 2, "available": 2, "image": "api:v2.1"}, + {"name": "worker", "replicas": 1, "available": 0, "image": "worker:v1.0"}, + ], + "production": [ + {"name": "app", "replicas": 2, "available": 2, "image": "app:v3.2"}, + {"name": "db", "replicas": 1, "available": 0, "image": "postgres:16"}, + ], +} + +SERVICES = { + "default": [ + {"name": "web-svc", "type": "LoadBalancer", "clusterIP": "10.0.0.10", "ports": "80/TCP"}, + {"name": "api-svc", "type": "ClusterIP", "clusterIP": "10.0.0.20", "ports": "8080/TCP"}, + ], + "production": [ + {"name": "app-svc", "type": "LoadBalancer", "clusterIP": "10.0.1.10", "ports": "443/TCP"}, + ], +} + +CONFIGMAPS = { + "default": [{"name": "app-config", "data_keys": ["DATABASE_URL", "LOG_LEVEL", "CACHE_TTL"]}], + "production": [{"name": "prod-config", "data_keys": ["DATABASE_URL", "REDIS_URL"]}], +} + +EVENTS = [ + { + "namespace": "default", + "type": "Warning", + "reason": "BackOff", + "object": "pod/worker-ghi56", + "message": "Back-off restarting failed container", + }, + { + "namespace": "production", + "type": "Warning", + "reason": "FailedScheduling", + "object": "pod/db-ggg77", + "message": "Insufficient memory on available nodes", + }, + { + "namespace": "default", + "type": "Normal", + "reason": "Pulled", + "object": "pod/api-def34", + "message": "Successfully pulled image api:v2.1", + }, + { + "namespace": "monitoring", + "type": "Normal", + "reason": "Started", + "object": "pod/prometheus-ccc33", + "message": "Started container prometheus", + }, +] + +LOGS = { + "web-abc12": ("2024-01-15T10:00:01Z GET /health 200 1ms\n2024-01-15T10:00:02Z GET /api/users 200 45ms\n"), + "api-def34": ( + "2024-01-15T10:00:01Z INFO Starting API server on :8080\n" + "2024-01-15T10:00:02Z WARN High latency detected: 250ms\n" + ), + "worker-ghi56": ( + "2024-01-15T10:00:01Z ERROR Connection refused: redis://redis:6379\n" + "2024-01-15T10:00:02Z FATAL Exiting due to unrecoverable error\n" + ), +} + +# Track mutable state for scale operations +_deployment_state: dict[str, dict[str, int]] = {} + + +# ============================================================================= +# Tool callbacks — each receives (params: dict, stdin: str | None) -> str +# ============================================================================= + + +def get_nodes(params, stdin=None): + """Return cluster nodes.""" + return json.dumps({"items": NODES}) + "\n" + + +def get_namespaces(params, stdin=None): + """Return namespaces.""" + return json.dumps({"items": NAMESPACES}) + "\n" + + +def get_pods(params, stdin=None): + """Return pods in namespace.""" + ns = params.get("namespace", "default") + pods = PODS.get(ns, []) + return json.dumps({"items": pods}) + "\n" + + +def get_deployments(params, stdin=None): + """Return deployments in namespace.""" + ns = params.get("namespace", "default") + deps = DEPLOYMENTS.get(ns, []) + return json.dumps({"items": deps}) + "\n" + + +def get_services(params, stdin=None): + """Return services in namespace.""" + ns = params.get("namespace", "default") + svcs = SERVICES.get(ns, []) + return json.dumps({"items": svcs}) + "\n" + + +def describe_pod(params, stdin=None): + """Describe a specific pod.""" + name = params.get("name", "") + ns = params.get("namespace", "default") + for pod in PODS.get(ns, []): + if pod["name"] == name: + detail = {**pod, "namespace": ns, "labels": {"app": name.rsplit("-", 1)[0]}} + return json.dumps(detail) + "\n" + raise ValueError(f"pod {name} not found in {ns}") + + +def get_logs(params, stdin=None): + """Get pod logs.""" + name = params.get("name", "") + tail = params.get("tail", 50) + logs = LOGS.get(name, f"No logs available for {name}\n") + lines = logs.strip().split("\n") + return "\n".join(lines[-int(tail) :]) + "\n" + + +def get_configmaps(params, stdin=None): + """List configmaps in namespace.""" + ns = params.get("namespace", "default") + cms = CONFIGMAPS.get(ns, []) + return json.dumps({"items": cms}) + "\n" + + +def get_secrets(params, stdin=None): + """List secrets (redacted).""" + ns = params.get("namespace", "default") + secrets = [{"name": f"{ns}-tls", "type": "kubernetes.io/tls", "data": "***REDACTED***"}] + return json.dumps({"items": secrets}) + "\n" + + +def get_events(params, stdin=None): + """Get cluster events, optionally filtered by namespace.""" + ns = params.get("namespace") + items = EVENTS if not ns else [e for e in EVENTS if e["namespace"] == ns] + return json.dumps({"items": items}) + "\n" + + +def scale_deployment(params, stdin=None): + """Scale a deployment.""" + name = params.get("name", "") + ns = params.get("namespace", "default") + replicas = params.get("replicas", 1) + key = f"{ns}/{name}" + _deployment_state[key] = {"replicas": int(replicas)} + return json.dumps({"deployment": name, "namespace": ns, "replicas": int(replicas), "status": "scaling"}) + "\n" + + +def rollout_status(params, stdin=None): + """Check rollout status of a deployment.""" + name = params.get("name", "") + ns = params.get("namespace", "default") + key = f"{ns}/{name}" + if key in _deployment_state: + r = _deployment_state[key]["replicas"] + return json.dumps({"deployment": name, "status": "progressing", "replicas": r, "updated": r}) + "\n" + for dep in DEPLOYMENTS.get(ns, []): + if dep["name"] == name: + status = "available" if dep["available"] == dep["replicas"] else "progressing" + return json.dumps({"deployment": name, "status": status, **dep}) + "\n" + raise ValueError(f"deployment {name} not found in {ns}") + + +# ============================================================================= +# Build the ScriptedTool with all 12 k8s commands +# ============================================================================= + + +def build_k8s_tool() -> ScriptedTool: + tool = ScriptedTool("kubectl", short_description="Kubernetes cluster management API") + + tool.add_tool("get_nodes", "List cluster nodes", callback=get_nodes) + + tool.add_tool("get_namespaces", "List namespaces", callback=get_namespaces) + + tool.add_tool( + "get_pods", + "List pods in a namespace", + callback=get_pods, + schema={"type": "object", "properties": {"namespace": {"type": "string", "description": "Namespace"}}}, + ) + + tool.add_tool( + "get_deployments", + "List deployments in a namespace", + callback=get_deployments, + schema={"type": "object", "properties": {"namespace": {"type": "string", "description": "Namespace"}}}, + ) + + tool.add_tool( + "get_services", + "List services in a namespace", + callback=get_services, + schema={"type": "object", "properties": {"namespace": {"type": "string", "description": "Namespace"}}}, + ) + + tool.add_tool( + "describe_pod", + "Describe a specific pod", + callback=describe_pod, + schema={ + "type": "object", + "properties": { + "name": {"type": "string", "description": "Pod name"}, + "namespace": {"type": "string", "description": "Namespace"}, + }, + "required": ["name"], + }, + ) + + tool.add_tool( + "get_logs", + "Get pod logs", + callback=get_logs, + schema={ + "type": "object", + "properties": { + "name": {"type": "string", "description": "Pod name"}, + "tail": {"type": "integer", "description": "Number of lines"}, + }, + "required": ["name"], + }, + ) + + tool.add_tool( + "get_configmaps", + "List configmaps in a namespace", + callback=get_configmaps, + schema={"type": "object", "properties": {"namespace": {"type": "string", "description": "Namespace"}}}, + ) + + tool.add_tool( + "get_secrets", + "List secrets in a namespace (values redacted)", + callback=get_secrets, + schema={"type": "object", "properties": {"namespace": {"type": "string", "description": "Namespace"}}}, + ) + + tool.add_tool( + "get_events", + "Get cluster events", + callback=get_events, + schema={"type": "object", "properties": {"namespace": {"type": "string", "description": "Filter namespace"}}}, + ) + + tool.add_tool( + "scale_deployment", + "Scale a deployment to N replicas", + callback=scale_deployment, + schema={ + "type": "object", + "properties": { + "name": {"type": "string", "description": "Deployment name"}, + "namespace": {"type": "string", "description": "Namespace"}, + "replicas": {"type": "integer", "description": "Target replica count"}, + }, + "required": ["name", "replicas"], + }, + ) + + tool.add_tool( + "rollout_status", + "Check deployment rollout status", + callback=rollout_status, + schema={ + "type": "object", + "properties": { + "name": {"type": "string", "description": "Deployment name"}, + "namespace": {"type": "string", "description": "Namespace"}, + }, + "required": ["name"], + }, + ) + + tool.env("KUBECONFIG", "/etc/kubernetes/admin.conf") + return tool + + +# ============================================================================= +# Demo scripts — what an LLM agent would generate +# ============================================================================= + + +def run_demos(tool: ScriptedTool) -> None: + print("=" * 70) + print("Kubernetes Orchestrator - 12 tools via ScriptedTool") + print("=" * 70) + + # -- Demo 1: Simple listing -- + print("\n--- Demo 1: List all nodes ---") + r = tool.execute_sync("get_nodes | jq -r '.items[] | \"\\(.name) \\(.status) cpu=\\(.cpu) mem=\\(.memory)\"'") + print(r.stdout) + + # -- Demo 2: Unhealthy pods across all namespaces -- + print("--- Demo 2: Find unhealthy pods across namespaces ---") + r = tool.execute_sync(""" + get_namespaces | jq -r '.items[].name' | while read ns; do + get_pods --namespace "$ns" \ + | jq -r '.items[] | select(.status != "Running") | .name + " " + .status' \ + | while read line; do echo " $ns/$line"; done + done + """) + print(r.stdout) + + # -- Demo 3: Cluster health report -- + print("--- Demo 3: Full cluster health report ---") + r = tool.execute_sync(""" + echo "=== Cluster Health Report ===" + + # Node status + echo "" + echo "-- Nodes --" + nodes=$(get_nodes) + total=$(echo "$nodes" | jq '.items | length') + ready=$(echo "$nodes" | jq '[.items[] | select(.status == "Ready")] | length') + echo "Nodes: $ready/$total ready" + + # Pod status per namespace + echo "" + echo "-- Pods --" + get_namespaces | jq -r '.items[].name' | while read ns; do + pods=$(get_pods --namespace "$ns") + total=$(echo "$pods" | jq '.items | length') + running=$(echo "$pods" | jq '[.items[] | select(.status == "Running")] | length') + echo " $ns: $running/$total running" + done + + # Warnings + echo "" + echo "-- Recent warnings --" + get_events | jq -r '.items[] | select(.type == "Warning") | " [\\(.reason)] \\(.object): \\(.message)"' + """) + print(r.stdout) + + # -- Demo 4: Diagnose CrashLoopBackOff -- + print("--- Demo 4: Diagnose crashing pod ---") + r = tool.execute_sync(""" + # Find pods in CrashLoopBackOff + crash_pods=$(get_pods --namespace default | jq -r '.items[] | select(.status == "CrashLoopBackOff") | .name') + + for pod in $crash_pods; do + echo "=== Diagnosing: $pod ===" + describe_pod --name "$pod" --namespace default | jq '{name, status, restarts, image, node}' + echo "" + echo "Recent logs:" + get_logs --name "$pod" --tail 5 + echo "Related events:" + get_events --namespace default | jq -r '.items[] | " [" + .type + "] " + .reason + ": " + .message' + echo "" + done + """) + print(r.stdout) + + # -- Demo 5: Scale + rollout -- + print("--- Demo 5: Scale deployment and check rollout ---") + r = tool.execute_sync(""" + echo "Scaling 'app' in production to 5 replicas..." + scale_deployment --name app --namespace production --replicas 5 | jq '.' + echo "" + echo "Rollout status:" + rollout_status --name app --namespace production | jq '.' + """) + print(r.stdout) + + # -- Demo 6: Service + configmap inventory -- + print("--- Demo 6: Namespace inventory ---") + r = tool.execute_sync(""" + for ns in default production; do + echo "=== Namespace: $ns ===" + echo "Services:" + get_services --namespace "$ns" | jq -r '.items[] | " \\(.name) (\\(.type)) -> \\(.ports)"' + echo "ConfigMaps:" + get_configmaps --namespace "$ns" | jq -r '.items[] | " \\(.name): \\(.data_keys | join(", "))"' + echo "Secrets:" + get_secrets --namespace "$ns" | jq -r '.items[] | " \\(.name) (\\(.type))"' + echo "" + done + """) + print(r.stdout) + + +# ============================================================================= +# LangChain integration (optional) +# ============================================================================= + + +def run_langchain_demo(tool: ScriptedTool) -> None: + """Demo: wrap ScriptedTool as LangChain tool for a ReAct agent.""" + try: + from langchain_anthropic import ChatAnthropic + from langgraph.prebuilt import create_react_agent + + from bashkit.langchain import create_scripted_tool + except ImportError as e: + print(f"\nSkipping LangChain demo (missing dependency: {e})") + print("Install with: pip install 'bashkit[langchain]' langgraph") + return + + print("\n" + "=" * 70) + print("LangChain ReAct Agent Demo") + print("=" * 70) + + # Wrap our k8s ScriptedTool as a LangChain tool + lc_tool = create_scripted_tool(tool) + print(f"\nLangChain tool: name={lc_tool.name!r}, tools={tool.tool_count()}") + + # Create agent with Claude + model = ChatAnthropic(model="claude-sonnet-4-20250514", temperature=0) + agent = create_react_agent(model, [lc_tool]) + + # Ask the agent to investigate the cluster + query = "Check the cluster health. Find any pods that are not running and diagnose why they're failing." + print(f"\nUser: {query}\n") + + result = agent.invoke({"messages": [{"role": "user", "content": query}]}) + # Print the final assistant message + for msg in result["messages"]: + if hasattr(msg, "content") and msg.type == "ai" and msg.content: + print(f"Agent: {msg.content}") + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + tool = build_k8s_tool() + + # Show what the LLM sees + print(f"Tool: {tool.name} ({tool.tool_count()} commands)\n") + print("--- System prompt (sent to LLM) ---") + print(tool.system_prompt()) + + # Run direct demos + run_demos(tool) + + # LangChain demo if requested + if "--langchain" in sys.argv: + run_langchain_demo(tool) + + print("=" * 70) + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/crates/bashkit-python/src/lib.rs b/crates/bashkit-python/src/lib.rs index cb19c451..c634f1eb 100644 --- a/crates/bashkit-python/src/lib.rs +++ b/crates/bashkit-python/src/lib.rs @@ -1,15 +1,100 @@ //! Python bindings for Bashkit //! -//! Exposes the Bash interpreter as a Python class for use in AI agent frameworks. -//! Uses stateful execution - filesystem and variables persist between calls. +//! Exposes the Bash interpreter and ScriptedTool as Python classes for use in +//! AI agent frameworks. BashTool provides stateful execution (filesystem persists +//! between calls). ScriptedTool composes Python callbacks as bash builtins for +//! multi-tool orchestration in a single script. -use bashkit::{Bash, BashTool as RustBashTool, ExecutionLimits, Tool}; +use bashkit::tool::VERSION; +use bashkit::{ + Bash, BashTool as RustBashTool, ExecutionLimits, ScriptedTool as RustScriptedTool, Tool, + ToolArgs, ToolDef, ToolRequest, +}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList}; use pyo3_async_runtimes::tokio::future_into_py; use std::sync::Arc; use tokio::sync::Mutex; +// ============================================================================ +// JSON <-> Python helpers +// ============================================================================ + +/// Convert serde_json::Value → PyObject +fn json_to_py(py: Python<'_>, val: &serde_json::Value) -> PyResult { + match val { + serde_json::Value::Null => Ok(py.None()), + serde_json::Value::Bool(b) => Ok(b.into_pyobject(py)?.to_owned().into_any().unbind()), + serde_json::Value::Number(n) => { + if let Some(i) = n.as_i64() { + Ok(i.into_pyobject(py)?.into_any().unbind()) + } else if let Some(f) = n.as_f64() { + Ok(f.into_pyobject(py)?.into_any().unbind()) + } else { + Ok(py.None()) + } + } + serde_json::Value::String(s) => Ok(s.into_pyobject(py)?.into_any().unbind()), + serde_json::Value::Array(arr) => { + let items: Vec = arr + .iter() + .map(|v| json_to_py(py, v)) + .collect::>()?; + Ok(PyList::new(py, &items)?.into_any().unbind()) + } + serde_json::Value::Object(map) => { + let dict = PyDict::new(py); + for (k, v) in map { + dict.set_item(k, json_to_py(py, v)?)?; + } + Ok(dict.into_any().unbind()) + } + } +} + +/// Convert PyObject → serde_json::Value (for schema dicts) +#[allow(clippy::only_used_in_recursion)] +fn py_to_json(py: Python<'_>, obj: &Bound<'_, pyo3::PyAny>) -> PyResult { + if obj.is_none() { + return Ok(serde_json::Value::Null); + } + if let Ok(b) = obj.extract::() { + return Ok(serde_json::Value::Bool(b)); + } + if let Ok(i) = obj.extract::() { + return Ok(serde_json::json!(i)); + } + if let Ok(f) = obj.extract::() { + return Ok(serde_json::json!(f)); + } + if let Ok(s) = obj.extract::() { + return Ok(serde_json::Value::String(s)); + } + if let Ok(list) = obj.downcast::() { + let arr: Vec = list + .iter() + .map(|item| py_to_json(py, &item)) + .collect::>()?; + return Ok(serde_json::Value::Array(arr)); + } + if let Ok(dict) = obj.downcast::() { + let mut map = serde_json::Map::new(); + for (k, v) in dict.iter() { + let key: String = k.extract()?; + map.insert(key, py_to_json(py, &v)?); + } + return Ok(serde_json::Value::Object(map)); + } + // Fallback: str() + let s = obj.str()?.extract::()?; + Ok(serde_json::Value::String(s)) +} + +// ============================================================================ +// ExecResult +// ============================================================================ + /// Result from executing bash commands #[pyclass] #[derive(Clone)] @@ -48,9 +133,9 @@ impl ExecResult { } /// Return output as dict - fn to_dict(&self) -> pyo3::PyResult> { + fn to_dict(&self) -> pyo3::PyResult> { Python::with_gil(|py| { - let dict = pyo3::types::PyDict::new(py); + let dict = PyDict::new(py); dict.set_item("stdout", &self.stdout)?; dict.set_item("stderr", &self.stderr)?; dict.set_item("exit_code", self.exit_code)?; @@ -60,6 +145,10 @@ impl ExecResult { } } +// ============================================================================ +// BashTool — stateful interpreter +// ============================================================================ + /// Virtual bash interpreter for AI agents /// /// BashTool provides a safe execution environment for running bash commands @@ -68,7 +157,7 @@ impl ExecResult { /// /// Example: /// ```python -/// from bashkit_py import BashTool +/// from bashkit import BashTool /// /// tool = BashTool() /// result = await tool.execute("echo 'Hello, World!'") @@ -77,7 +166,6 @@ impl ExecResult { #[pyclass] #[allow(dead_code)] pub struct BashTool { - /// Stateful bash interpreter - persists filesystem and variables inner: Arc>, username: Option, hostname: Option, @@ -87,13 +175,6 @@ pub struct BashTool { #[pymethods] impl BashTool { - /// Create a new BashTool instance - /// - /// Args: - /// username: Custom username for virtual environment (default: "user") - /// hostname: Custom hostname for virtual environment (default: "sandbox") - /// max_commands: Maximum commands to execute (default: 10000) - /// max_loop_iterations: Maximum loop iterations (default: 100000) #[new] #[pyo3(signature = (username=None, hostname=None, max_commands=None, max_loop_iterations=None))] fn new( @@ -131,22 +212,6 @@ impl BashTool { }) } - /// Execute bash commands asynchronously - /// - /// State persists between calls - files, variables, and functions - /// created in one call are available in subsequent calls. - /// - /// Args: - /// commands: Bash commands to execute (like `bash -c "commands"`) - /// - /// Returns: - /// ExecResult with stdout, stderr, exit_code - /// - /// Example: - /// ```python - /// result = await tool.execute("echo hello && echo world") - /// print(result.stdout) # hello\nworld\n - /// ``` fn execute<'py>(&self, py: Python<'py>, commands: String) -> PyResult> { let inner = self.inner.clone(); future_into_py(py, async move { @@ -168,15 +233,6 @@ impl BashTool { }) } - /// Execute bash commands synchronously (blocking) - /// - /// Note: Prefer `execute()` for async contexts. This method blocks. - /// - /// Args: - /// commands: Bash commands to execute - /// - /// Returns: - /// ExecResult with stdout, stderr, exit_code fn execute_sync(&self, commands: String) -> PyResult { let inner = self.inner.clone(); let rt = tokio::runtime::Runtime::new() @@ -201,7 +257,6 @@ impl BashTool { }) } - /// Reset the interpreter state (clear filesystem, variables, functions) fn reset(&self) -> PyResult<()> { let inner = self.inner.clone(); let rt = tokio::runtime::Runtime::new() @@ -209,45 +264,37 @@ impl BashTool { rt.block_on(async move { let mut bash = inner.lock().await; - // Create fresh Bash with same settings let builder = Bash::builder(); - // Note: We lose settings on reset, could store them *bash = builder.build(); Ok(()) }) } - /// Get the tool name #[getter] fn name(&self) -> &str { "bashkit" } - /// Get short description #[getter] fn short_description(&self) -> &str { "Virtual bash interpreter with virtual filesystem" } - /// Get the full description fn description(&self) -> PyResult { let tool = RustBashTool::default(); Ok(tool.description()) } - /// Get LLM documentation fn help(&self) -> PyResult { let tool = RustBashTool::default(); Ok(tool.help()) } - /// Get system prompt for LLMs fn system_prompt(&self) -> PyResult { let tool = RustBashTool::default(); Ok(tool.system_prompt()) } - /// Get JSON schema for input validation fn input_schema(&self) -> PyResult { let tool = RustBashTool::default(); let schema = tool.input_schema(); @@ -255,7 +302,6 @@ impl BashTool { .map_err(|e| PyValueError::new_err(format!("Schema serialization failed: {}", e))) } - /// Get JSON schema for output fn output_schema(&self) -> PyResult { let tool = RustBashTool::default(); let schema = tool.output_schema(); @@ -263,10 +309,9 @@ impl BashTool { .map_err(|e| PyValueError::new_err(format!("Schema serialization failed: {}", e))) } - /// Get tool version #[getter] fn version(&self) -> &str { - bashkit::tool::VERSION + VERSION } fn __repr__(&self) -> String { @@ -278,26 +323,275 @@ impl BashTool { } } -/// Create a LangChain-compatible tool from BashTool +// ============================================================================ +// ScriptedTool — multi-tool orchestration via bash scripts +// ============================================================================ + +/// Entry for a registered Python tool callback +struct PyToolEntry { + name: String, + description: String, + schema: serde_json::Value, + callback: PyObject, +} + +/// Compose Python callbacks as bash builtins for multi-tool orchestration. /// -/// Returns a dict with: -/// - name: Tool name -/// - description: Tool description -/// - args_schema: JSON schema for arguments +/// Each registered tool becomes a bash builtin command. An LLM (or user) writes +/// a single bash script that pipes, loops, and branches across all tools. +/// +/// Python callbacks receive `(params: dict, stdin: str | None)` and return a +/// string. Raise an exception to signal failure. /// /// Example: /// ```python -/// from bashkit_py import create_langchain_tool_spec +/// from bashkit import ScriptedTool +/// +/// def get_user(params, stdin=None): +/// return '{"id": 1, "name": "Alice"}' /// -/// spec = create_langchain_tool_spec() -/// # Use with langchain's StructuredTool.from_function() +/// tool = ScriptedTool("api") +/// tool.add_tool("get_user", "Fetch user by ID", +/// callback=get_user, +/// schema={"type": "object", "properties": {"id": {"type": "integer"}}}) +/// +/// result = tool.execute_sync("get_user --id 1 | jq -r '.name'") +/// print(result.stdout) # Alice /// ``` +#[pyclass] +pub struct ScriptedTool { + name: String, + short_desc: Option, + tools: Vec, + env_vars: Vec<(String, String)>, + max_commands: Option, + max_loop_iterations: Option, +} + +impl ScriptedTool { + /// Build a Rust ScriptedTool from stored Python config. + /// Each Python callback is wrapped via `Python::with_gil`. + fn build_rust_tool(&self) -> RustScriptedTool { + let mut builder = RustScriptedTool::builder(&self.name); + + if let Some(ref desc) = self.short_desc { + builder = builder.short_description(desc); + } + + for entry in &self.tools { + let py_cb = Python::with_gil(|py| entry.callback.clone_ref(py)); + let tool_name = entry.name.clone(); + + let callback = move |args: &ToolArgs| -> Result { + Python::with_gil(|py| { + let params = json_to_py(py, &args.params).map_err(|e| e.to_string())?; + let stdin_arg = args.stdin.as_deref().map(|s| s.to_string()); + + let result = py_cb + .call1(py, (params, stdin_arg)) + .map_err(|e| format!("{}: {}", tool_name, e))?; + result + .extract::(py) + .map_err(|e| format!("{}: callback must return str, got {}", tool_name, e)) + }) + }; + + builder = builder.tool( + ToolDef::new(&entry.name, &entry.description).with_schema(entry.schema.clone()), + callback, + ); + } + + for (k, v) in &self.env_vars { + builder = builder.env(k, v); + } + + if self.max_commands.is_some() || self.max_loop_iterations.is_some() { + let mut limits = ExecutionLimits::new(); + if let Some(mc) = self.max_commands { + limits = limits.max_commands(mc as usize); + } + if let Some(mli) = self.max_loop_iterations { + limits = limits.max_loop_iterations(mli as usize); + } + builder = builder.limits(limits); + } + + builder.build() + } +} + +#[pymethods] +impl ScriptedTool { + /// Create a new ScriptedTool. + /// + /// Args: + /// name: Tool name (used in system prompt and docs) + /// short_description: One-line description + /// max_commands: Max commands per execute call + /// max_loop_iterations: Max loop iterations per execute call + #[new] + #[pyo3(signature = (name, short_description=None, max_commands=None, max_loop_iterations=None))] + fn new( + name: String, + short_description: Option, + max_commands: Option, + max_loop_iterations: Option, + ) -> Self { + Self { + name, + short_desc: short_description, + tools: Vec::new(), + env_vars: Vec::new(), + max_commands, + max_loop_iterations, + } + } + + /// Register a tool command. + /// + /// The callback signature is: `callback(params: dict, stdin: str | None) -> str` + /// + /// `params` contains `--key value` flags parsed from the bash command line, + /// with types coerced per the schema (integers, booleans, etc.). + /// + /// Args: + /// name: Command name (becomes a bash builtin) + /// description: Human-readable description + /// callback: Python callable `(params, stdin) -> str` + /// schema: Optional JSON Schema dict for input parameters + #[pyo3(signature = (name, description, callback, schema=None))] + fn add_tool( + &mut self, + py: Python<'_>, + name: String, + description: String, + callback: PyObject, + schema: Option>, + ) -> PyResult<()> { + let schema_val = match schema { + Some(ref s) => py_to_json(py, s)?, + None => serde_json::Value::Object(Default::default()), + }; + self.tools.push(PyToolEntry { + name, + description, + schema: schema_val, + callback, + }); + Ok(()) + } + + /// Add an environment variable visible inside scripts. + fn env(&mut self, key: String, value: String) { + self.env_vars.push((key, value)); + } + + /// Execute a bash script asynchronously. + fn execute<'py>(&self, py: Python<'py>, commands: String) -> PyResult> { + let mut tool = self.build_rust_tool(); + future_into_py(py, async move { + let resp = tool.execute(ToolRequest { commands }).await; + Ok(ExecResult { + stdout: resp.stdout, + stderr: resp.stderr, + exit_code: resp.exit_code, + error: resp.error, + }) + }) + } + + /// Execute a bash script synchronously (blocking). + fn execute_sync(&self, commands: String) -> PyResult { + let mut tool = self.build_rust_tool(); + let rt = tokio::runtime::Runtime::new() + .map_err(|e| PyRuntimeError::new_err(format!("Failed to create runtime: {}", e)))?; + + let resp = rt.block_on(async move { tool.execute(ToolRequest { commands }).await }); + Ok(ExecResult { + stdout: resp.stdout, + stderr: resp.stderr, + exit_code: resp.exit_code, + error: resp.error, + }) + } + + /// Get the tool name. + #[getter(name)] + fn name_prop(&self) -> &str { + &self.name + } + + /// Get the short description. + #[getter] + fn short_description(&self) -> String { + self.short_desc + .clone() + .unwrap_or_else(|| format!("ScriptedTool: {}", self.name)) + } + + /// Number of registered tools. + fn tool_count(&self) -> usize { + self.tools.len() + } + + /// Get the full description. + fn description(&self) -> String { + self.build_rust_tool().description() + } + + /// Get help text (man-page format). + fn help(&self) -> String { + self.build_rust_tool().help() + } + + /// Get system prompt for LLMs (token-efficient). + fn system_prompt(&self) -> String { + self.build_rust_tool().system_prompt() + } + + /// Get JSON input schema. + fn input_schema(&self) -> PyResult { + let tool = self.build_rust_tool(); + let schema = tool.input_schema(); + serde_json::to_string_pretty(&schema) + .map_err(|e| PyValueError::new_err(format!("Schema serialization failed: {}", e))) + } + + /// Get JSON output schema. + fn output_schema(&self) -> PyResult { + let tool = self.build_rust_tool(); + let schema = tool.output_schema(); + serde_json::to_string_pretty(&schema) + .map_err(|e| PyValueError::new_err(format!("Schema serialization failed: {}", e))) + } + + /// Get tool version. + #[getter] + fn version(&self) -> &str { + VERSION + } + + fn __repr__(&self) -> String { + format!( + "ScriptedTool(name={:?}, tools={})", + self.name, + self.tools.len() + ) + } +} + +// ============================================================================ +// Module-level functions +// ============================================================================ + +/// Create a LangChain-compatible tool spec from BashTool. #[pyfunction] -fn create_langchain_tool_spec() -> PyResult> { +fn create_langchain_tool_spec() -> PyResult> { let tool = RustBashTool::default(); Python::with_gil(|py| { - let dict = pyo3::types::PyDict::new(py); + let dict = PyDict::new(py); dict.set_item("name", tool.name())?; dict.set_item("description", tool.description())?; @@ -310,10 +604,14 @@ fn create_langchain_tool_spec() -> PyResult> { }) } -/// Python module definition +// ============================================================================ +// Python module +// ============================================================================ + #[pymodule] fn _bashkit(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_function(wrap_pyfunction!(create_langchain_tool_spec, m)?)?; Ok(()) diff --git a/crates/bashkit-python/tests/test_bashkit.py b/crates/bashkit-python/tests/test_bashkit.py index a1daa1e3..61496482 100644 --- a/crates/bashkit-python/tests/test_bashkit.py +++ b/crates/bashkit-python/tests/test_bashkit.py @@ -4,9 +4,9 @@ import pytest -from bashkit import BashTool, create_langchain_tool_spec +from bashkit import BashTool, ScriptedTool, create_langchain_tool_spec -# -- Construction ----------------------------------------------------------- +# -- BashTool: Construction ------------------------------------------------- def test_default_construction(): @@ -21,7 +21,7 @@ def test_custom_construction(): assert repr(tool) == 'BashTool(username="alice", hostname="box")' -# -- Sync execution --------------------------------------------------------- +# -- BashTool: Sync execution ----------------------------------------------- def test_echo(): @@ -71,7 +71,7 @@ def test_file_persistence(): assert r.stdout.strip() == "content" -# -- Async execution -------------------------------------------------------- +# -- BashTool: Async execution ---------------------------------------------- @pytest.mark.asyncio @@ -121,7 +121,7 @@ def test_exec_result_str_failure(): assert "Error" in str(r) -# -- Reset ------------------------------------------------------------------ +# -- BashTool: Reset -------------------------------------------------------- def test_reset(): @@ -132,7 +132,7 @@ def test_reset(): assert r.stdout.strip() == "empty" -# -- LLM metadata ---------------------------------------------------------- +# -- BashTool: LLM metadata ------------------------------------------------ def test_description(): @@ -179,3 +179,309 @@ def test_langchain_tool_spec(): assert "description" in spec assert "args_schema" in spec assert spec["name"] == "bashkit" + + +# =========================================================================== +# ScriptedTool tests +# =========================================================================== + + +def _make_echo_tool(): + """Helper: ScriptedTool with one 'greet' command.""" + tool = ScriptedTool("test_api", short_description="Test API") + tool.add_tool( + "greet", + "Greet a user", + callback=lambda params, stdin=None: f"hello {params.get('name', 'world')}\n", + schema={"type": "object", "properties": {"name": {"type": "string"}}}, + ) + return tool + + +# -- ScriptedTool: Construction --------------------------------------------- + + +def test_scripted_tool_construction(): + tool = ScriptedTool("my_api") + assert tool.name == "my_api" + assert tool.tool_count() == 0 + assert "ScriptedTool" in tool.short_description + + +def test_scripted_tool_custom_description(): + tool = ScriptedTool("api", short_description="My custom API") + assert tool.short_description == "My custom API" + + +def test_scripted_tool_repr(): + tool = _make_echo_tool() + assert "test_api" in repr(tool) + assert "1" in repr(tool) # tool count + + +# -- ScriptedTool: add_tool ------------------------------------------------- + + +def test_add_tool_increments_count(): + tool = ScriptedTool("api") + assert tool.tool_count() == 0 + tool.add_tool("cmd1", "Command 1", callback=lambda p, s=None: "ok\n") + assert tool.tool_count() == 1 + tool.add_tool("cmd2", "Command 2", callback=lambda p, s=None: "ok\n") + assert tool.tool_count() == 2 + + +def test_add_tool_with_schema(): + tool = ScriptedTool("api") + tool.add_tool( + "get_user", + "Fetch user", + callback=lambda p, s=None: json.dumps({"id": p.get("id", 0)}) + "\n", + schema={"type": "object", "properties": {"id": {"type": "integer"}}}, + ) + assert tool.tool_count() == 1 + + +def test_add_tool_no_schema(): + tool = ScriptedTool("api") + tool.add_tool("noop", "No-op", callback=lambda p, s=None: "ok\n") + assert tool.tool_count() == 1 + + +# -- ScriptedTool: execute_sync -------------------------------------------- + + +def test_scripted_tool_single_call(): + tool = _make_echo_tool() + r = tool.execute_sync("greet --name Alice") + assert r.exit_code == 0 + assert r.stdout.strip() == "hello Alice" + + +def test_scripted_tool_pipeline_with_jq(): + tool = ScriptedTool("api") + tool.add_tool( + "get_user", + "Fetch user", + callback=lambda p, s=None: '{"id": 1, "name": "Alice"}\n', + ) + r = tool.execute_sync("get_user | jq -r '.name'") + assert r.exit_code == 0 + assert r.stdout.strip() == "Alice" + + +def test_scripted_tool_multi_step(): + tool = ScriptedTool("api") + tool.add_tool( + "get_user", + "Fetch user", + callback=lambda p, s=None: f'{{"id": {p.get("id", 0)}, "name": "Bob"}}\n', + schema={"type": "object", "properties": {"id": {"type": "integer"}}}, + ) + tool.add_tool( + "get_orders", + "Fetch orders", + callback=lambda p, s=None: '[{"total": 10}, {"total": 20}]\n', + schema={"type": "object", "properties": {"user_id": {"type": "integer"}}}, + ) + r = tool.execute_sync(""" + user=$(get_user --id 1) + name=$(echo "$user" | jq -r '.name') + total=$(get_orders --user_id 1 | jq '[.[].total] | add') + echo "$name: $total" + """) + assert r.exit_code == 0 + assert r.stdout.strip() == "Bob: 30" + + +def test_scripted_tool_callback_error(): + tool = ScriptedTool("api") + tool.add_tool( + "fail_cmd", + "Always fails", + callback=lambda p, s=None: (_ for _ in ()).throw(ValueError("service down")), + ) + r = tool.execute_sync("fail_cmd") + assert r.exit_code != 0 + assert "service down" in r.stderr + + +def test_scripted_tool_error_fallback(): + tool = ScriptedTool("api") + tool.add_tool( + "fail_cmd", + "Always fails", + callback=lambda p, s=None: (_ for _ in ()).throw(ValueError("boom")), + ) + r = tool.execute_sync("fail_cmd || echo fallback") + assert r.exit_code == 0 + assert "fallback" in r.stdout + + +def test_scripted_tool_stdin_pipe(): + tool = ScriptedTool("api") + tool.add_tool( + "upper", + "Uppercase stdin", + callback=lambda p, stdin=None: (stdin or "").upper(), + ) + r = tool.execute_sync("echo hello | upper") + assert r.exit_code == 0 + assert r.stdout.strip() == "HELLO" + + +def test_scripted_tool_env_var(): + tool = ScriptedTool("api") + tool.env("API_URL", "https://example.com") + tool.add_tool("noop", "No-op", callback=lambda p, s=None: "ok\n") + r = tool.execute_sync("echo $API_URL") + assert r.exit_code == 0 + assert r.stdout.strip() == "https://example.com" + + +def test_scripted_tool_loop(): + tool = ScriptedTool("api") + tool.add_tool( + "get_user", + "Fetch user", + callback=lambda p, s=None: f'{{"name": "user{p.get("id", 0)}"}}\n', + schema={"type": "object", "properties": {"id": {"type": "integer"}}}, + ) + r = tool.execute_sync(""" + for uid in 1 2 3; do + get_user --id $uid | jq -r '.name' + done + """) + assert r.exit_code == 0 + assert r.stdout.strip() == "user1\nuser2\nuser3" + + +def test_scripted_tool_conditional(): + tool = ScriptedTool("api") + tool.add_tool( + "check", + "Check status", + callback=lambda p, s=None: '{"ok": true}\n', + ) + r = tool.execute_sync(""" + status=$(check | jq -r '.ok') + if [ "$status" = "true" ]; then + echo "healthy" + else + echo "unhealthy" + fi + """) + assert r.exit_code == 0 + assert r.stdout.strip() == "healthy" + + +def test_scripted_tool_multiple_execute(): + """Multiple execute calls on the same tool work (stateless between calls).""" + tool = _make_echo_tool() + r1 = tool.execute_sync("greet --name Alice") + assert r1.stdout.strip() == "hello Alice" + r2 = tool.execute_sync("greet --name Bob") + assert r2.stdout.strip() == "hello Bob" + + +def test_scripted_tool_empty_script(): + tool = _make_echo_tool() + r = tool.execute_sync("") + assert r.exit_code == 0 + assert r.stdout == "" + + +def test_scripted_tool_boolean_flag(): + tool = ScriptedTool("api") + tool.add_tool( + "search", + "Search", + callback=lambda p, s=None: f"verbose={p.get('verbose', False)}\n", + schema={"type": "object", "properties": {"verbose": {"type": "boolean"}}}, + ) + r = tool.execute_sync("search --verbose") + assert r.exit_code == 0 + assert r.stdout.strip() == "verbose=True" + + +def test_scripted_tool_integer_coercion(): + tool = ScriptedTool("api") + tool.add_tool( + "get", + "Get by ID", + callback=lambda p, s=None: f"id={p.get('id')} type={type(p.get('id')).__name__}\n", + schema={"type": "object", "properties": {"id": {"type": "integer"}}}, + ) + r = tool.execute_sync("get --id 42") + assert r.exit_code == 0 + assert r.stdout.strip() == "id=42 type=int" + + +# -- ScriptedTool: Async execution ----------------------------------------- + + +@pytest.mark.asyncio +async def test_scripted_tool_async_execute(): + tool = _make_echo_tool() + r = await tool.execute("greet --name Async") + assert r.exit_code == 0 + assert r.stdout.strip() == "hello Async" + + +# -- ScriptedTool: Introspection ------------------------------------------- + + +def test_scripted_tool_system_prompt(): + tool = _make_echo_tool() + sp = tool.system_prompt() + assert "# test_api" in sp + assert "greet" in sp + assert "--name" in sp + + +def test_scripted_tool_description(): + tool = _make_echo_tool() + desc = tool.description() + assert "greet" in desc + + +def test_scripted_tool_help(): + tool = _make_echo_tool() + h = tool.help() + assert "TOOL COMMANDS" in h + assert "greet" in h + + +def test_scripted_tool_schemas(): + tool = _make_echo_tool() + inp = json.loads(tool.input_schema()) + assert "properties" in inp + out = json.loads(tool.output_schema()) + assert "properties" in out + + +def test_scripted_tool_version(): + tool = _make_echo_tool() + assert isinstance(tool.version, str) + assert len(tool.version) > 0 + + +# -- ScriptedTool: Many tools (12) ----------------------------------------- + + +def test_scripted_tool_dozen_tools(): + """Register 12 tools and execute a multi-tool script.""" + tool = ScriptedTool("big_api", short_description="API with 12 commands") + for i in range(12): + name = f"cmd{i}" + tool.add_tool( + name, + f"Command {i}", + callback=lambda p, s=None, idx=i: f"result-{idx}\n", + ) + assert tool.tool_count() == 12 + # Call all 12 + r = tool.execute_sync("; ".join(f"cmd{i}" for i in range(12))) + assert r.exit_code == 0 + lines = r.stdout.strip().splitlines() + assert lines == [f"result-{i}" for i in range(12)]