@datalayer/agent-runtimes 1.0.4 → 1.0.6

package/scripts/codegen/generate_agents.py

@@ -11,6 +11,7 @@ Generates Python and TypeScript code from YAML agent specifications.
 import argparse
 import json
 import sys
+from copy import deepcopy
 from pathlib import Path
 from typing import Any, Dict, List
 
@@ -27,7 +28,7 @@ def _fmt_list(items: list[str]) -> str:
     """Format a list of strings with double quotes for ruff compliance."""
     if not items:
         return "[]"
-    return "[" + ", ".join(f'"{item}"' for item in items) + "]"
+    return "[" + ", ".join(_fmt_py_literal(item) for item in items) + "]"
 
 
 def _fmt_py_literal(value: Any) -> str:
@@ -44,6 +45,68 @@ def _fmt_ts_literal(value: Any) -> str:
     return json.dumps(value, ensure_ascii=False)
 
 
+def _normalize_subagents_for_typescript(value: Any) -> Any:
+    """Convert subagents config keys from YAML snake_case to TS camelCase."""
+    if not isinstance(value, dict):
+        return value
+
+    top_level_map = {
+        "default_model": "defaultModel",
+        "include_general_purpose": "includeGeneralPurpose",
+        "max_nesting_depth": "maxNestingDepth",
+    }
+    subagent_map = {
+        "can_ask_questions": "canAskQuestions",
+        "max_questions": "maxQuestions",
+        "preferred_mode": "preferredMode",
+        "typical_complexity": "typicalComplexity",
+        "typically_needs_context": "typicallyNeedsContext",
+    }
+
+    normalized: dict[str, Any] = {}
+    for key, raw_val in value.items():
+        mapped_key = top_level_map.get(key, key)
+        if mapped_key == "subagents" and isinstance(raw_val, list):
+            normalized_subagents: list[Any] = []
+            for subagent in raw_val:
+                if isinstance(subagent, dict):
+                    normalized_subagent: dict[str, Any] = {}
+                    for sa_key, sa_val in subagent.items():
+                        normalized_subagent[subagent_map.get(sa_key, sa_key)] = sa_val
+                    normalized_subagents.append(normalized_subagent)
+                else:
+                    normalized_subagents.append(subagent)
+            normalized[mapped_key] = normalized_subagents
+        else:
+            normalized[mapped_key] = raw_val
+
+    return normalized
+
+
+def _sanitize_tool_hooks_for_codegen(value: Any) -> Any:
+    """Normalize tool hook paths that would trigger bandit in generated code."""
+    if not isinstance(value, dict):
+        return value
+
+    tool_hooks = deepcopy(value)
+    for key in ("audit_log_path", "auditLogPath"):
+        raw_path = tool_hooks.get(key)
+        if isinstance(raw_path, str):
+            raw_parts = Path(raw_path).parts
+            if len(raw_parts) >= 3 and raw_parts[0] == "/" and raw_parts[1] == "tmp":
+                tool_hooks[key] = Path(raw_path).name
+    return tool_hooks
+
+
+def _sanitize_spec_for_codegen(spec: Dict[str, Any]) -> Dict[str, Any]:
+    """Return a spec copy sanitized for generated Python/TypeScript outputs."""
+    sanitized = deepcopy(spec)
+    sanitized["tool_hooks"] = _sanitize_tool_hooks_for_codegen(
+        sanitized.get("tool_hooks")
+    )
+    return sanitized
+
+
 def load_yaml_specs(specs_dir: Path) -> List[tuple[str, Dict[str, Any]]]:
     """
     Load all YAML agent specifications from directory and subdirectories.
@@ -91,7 +154,7 @@ Generated from YAML specifications in specs/agents/
 from typing import Dict
 
 from agent_runtimes.mcp.catalog_mcp_servers import MCP_SERVER_CATALOG
-from agent_runtimes.types import AgentSpec
+from agent_runtimes.types import AgentSpec, SubAgentSpecConfig, SubAgentsConfig
 
 # ============================================================================
 # Agent Specs
@@ -123,6 +186,7 @@ from agent_runtimes.types import AgentSpec
             code += f"# {'=' * 76}\n\n"
 
         for spec in folder_specs:
+            spec = _sanitize_spec_for_codegen(spec)
             agent_id = spec["id"]
             version = spec["version"]
             # Prefix agent ID with folder name for uniqueness
@@ -166,7 +230,7 @@ from agent_runtimes.types import AgentSpec
             suggestions = spec.get("suggestions", [])
             suggestions_str = (
                 "[\n        "
-                + ",\n        ".join(f'"{s}"' for s in suggestions)
+                + ",\n        ".join(_fmt_py_literal(s) for s in suggestions)
                 + ",\n    ]"
                 if suggestions
                 else "[]"
@@ -210,6 +274,10 @@ from agent_runtimes.types import AgentSpec
             # Model field
             model_id = spec.get("model")
             model_str = f'"{model_id}"' if model_id else "None"
+            inference_provider = spec.get("inference_provider")
+            inference_provider_str = (
+                f'"{inference_provider}"' if inference_provider else "None"
+            )
 
             # Sandbox variant field
             sandbox_variant = spec.get("sandbox_variant")
@@ -217,11 +285,12 @@ from agent_runtimes.types import AgentSpec
 
             # New flow-level fields
             goal_raw = spec.get("goal")
-            goal_str = (
-                f'"{goal_raw.replace(chr(10), " ").replace("  ", " ").strip()}"'
+            goal_clean = (
+                goal_raw.replace(chr(10), " ").replace("  ", " ").strip()
                 if goal_raw
-                else "None"
+                else None
             )
+            goal_str = _fmt_py_literal(goal_clean)
             protocol_val = spec.get("protocol")
             protocol_str = f'"{protocol_val}"' if protocol_val else "None"
             ui_ext = spec.get("ui_extension")
@@ -239,6 +308,49 @@ from agent_runtimes.types import AgentSpec
             notifs = spec.get("notifications")
             memory_val = spec.get("memory")
             memory_str = f'"{memory_val}"' if memory_val else "None"
+            pre_hooks_val = spec.get("pre_hooks")
+            post_hooks_val = spec.get("post_hooks")
+            tool_hooks_val = spec.get("tool_hooks")
+            parameters_val = spec.get("parameters")
+            subagents_val = spec.get("subagents")
+
+            # Build subagents code if present
+            subagents_str = "None"
+            if isinstance(subagents_val, dict) and subagents_val.get("subagents"):
+                sa_items = []
+                for sa in subagents_val["subagents"]:
+                    sa_fields = [
+                        f"name={_fmt_py_literal(sa['name'])}",
+                        f"description={_fmt_py_literal(sa['description'])}",
+                        f"instructions={_fmt_py_literal(sa['instructions'])}",
+                    ]
+                    for opt_key in (
+                        "model",
+                        "can_ask_questions",
+                        "max_questions",
+                        "preferred_mode",
+                        "typical_complexity",
+                        "typically_needs_context",
+                    ):
+                        opt_val = sa.get(opt_key)
+                        if opt_val is not None:
+                            sa_fields.append(f"{opt_key}={_fmt_py_literal(opt_val)}")
+                    sa_items.append("SubAgentSpecConfig(" + ", ".join(sa_fields) + ")")
+                sa_list_str = "[" + ", ".join(sa_items) + "]"
+                cfg_parts = [f"subagents={sa_list_str}"]
+                if subagents_val.get("default_model") is not None:
+                    cfg_parts.append(
+                        f"default_model={_fmt_py_literal(subagents_val['default_model'])}"
+                    )
+                if subagents_val.get("include_general_purpose") is not None:
+                    cfg_parts.append(
+                        f"include_general_purpose={_fmt_py_literal(subagents_val['include_general_purpose'])}"
+                    )
+                if subagents_val.get("max_nesting_depth") is not None:
+                    cfg_parts.append(
+                        f"max_nesting_depth={_fmt_py_literal(subagents_val['max_nesting_depth'])}"
+                    )
+                subagents_str = "SubAgentsConfig(" + ", ".join(cfg_parts) + ")"
 
             code += f'''{const_name} = AgentSpec(
     id="{full_agent_id}",
@@ -248,6 +360,7 @@ from agent_runtimes.types import AgentSpec
     tags={_fmt_list(spec.get("tags", []))},
     enabled={spec.get("enabled", True)},
     model={model_str},
+    inference_provider={inference_provider_str},
     mcp_servers=[{mcp_servers_str}],
     skills={_fmt_list(skill_refs)},
     tools={_fmt_list(tool_refs)},
@@ -277,6 +390,11 @@ from agent_runtimes.types import AgentSpec
     authorization_policy={auth_policy_str},
     notifications={_fmt_py_literal(notifs)},
     memory={memory_str},
+    pre_hooks={_fmt_py_literal(pre_hooks_val)},
+    post_hooks={_fmt_py_literal(post_hooks_val)},
+    tool_hooks={_fmt_py_literal(tool_hooks_val)},
+    parameters={_fmt_py_literal(parameters_val)},
+    subagents={subagents_str},
 )
 
 '''
@@ -564,15 +682,15 @@ const SKILL_MAP: Record<string, any> = {
         code += "};\n"
         code += """
 function toAgentSkillSpec(skill: SkillSpec) {
-  return {
-    id: skill.id,
-    name: skill.name,
-    description: skill.description,
+    return {
+        id: skill.id,
+        name: skill.name,
+        description: skill.description,
         version: skill.version ?? '0.0.1',
-    tags: skill.tags,
-    enabled: skill.enabled,
-    requiredEnvVars: skill.requiredEnvVars,
-  };
+        tags: skill.tags,
+        enabled: skill.enabled,
+        requiredEnvVars: skill.requiredEnvVars,
+    };
 }
 """
 
@@ -629,6 +747,7 @@ const FRONTEND_TOOL_MAP: Record<string, any> = {
             code += f"// {'=' * 76}\n\n"
 
         for spec in folder_specs:
+            spec = _sanitize_spec_for_codegen(spec)
             agent_id = spec["id"]
             version = spec["version"]
             # Prefix agent ID with folder name for uniqueness
@@ -663,7 +782,8 @@ const FRONTEND_TOOL_MAP: Record<string, any> = {
             ]
             if has_skills and skill_ids_list:
                 skills_str = ", ".join(
-                    f"toAgentSkillSpec(SKILL_MAP['{sid}'])" for sid in skill_ids_list
+                    f"(SKILL_MAP['{sid}'] ? toAgentSkillSpec(SKILL_MAP['{sid}']) : undefined)"
+                    for sid in skill_ids_list
                 )
             else:
                 skills_str = ""
@@ -731,6 +851,10 @@ const FRONTEND_TOOL_MAP: Record<string, any> = {
             # Model field
             model_id = spec.get("model")
             model_ts = f"'{model_id}'" if model_id else "undefined"
+            inference_provider = spec.get("inference_provider")
+            inference_provider_ts = (
+                f"'{inference_provider}'" if inference_provider else "undefined"
+            )
 
             # Sandbox variant field
             sandbox_variant = spec.get("sandbox_variant")
@@ -764,44 +888,62 @@ const FRONTEND_TOOL_MAP: Record<string, any> = {
             notifs = spec.get("notifications")
             memory_val = spec.get("memory")
             memory_ts = f"'{memory_val}'" if memory_val else "undefined"
+            pre_hooks_val = spec.get("pre_hooks")
+            post_hooks_val = spec.get("post_hooks")
+            tool_hooks_val = spec.get("tool_hooks")
+            parameters_val = spec.get("parameters")
+            subagents_val = spec.get("subagents")
+            subagents_ts = _fmt_ts_literal(
+                _normalize_subagents_for_typescript(subagents_val)
+            )
+            inference_provider_line = (
+                f"    inferenceProvider: {inference_provider_ts},\n"
+                if inference_provider
+                else ""
+            )
 
             code += f"""export const {const_name}: AgentSpec = {{
-  id: '{full_agent_id}',
+    id: '{full_agent_id}',
     version: '{version}',
-  name: '{spec["name"]}',
-  description: `{description}`,
-  tags: {tags_str},
-  enabled: {str(spec.get("enabled", True)).lower()},
-  model: {model_ts},
-  mcpServers: [{mcp_servers_str}],
-  skills: [{skills_str}],
+    name: '{spec["name"]}',
+    description: `{description}`,
+    tags: {tags_str},
+    enabled: {str(spec.get("enabled", True)).lower()},
+    model: {model_ts},
+{inference_provider_line}    mcpServers: [{mcp_servers_str}],
+    skills: [{skills_str}].filter(Boolean) as SkillSpec[],
     tools: [{tools_str}],
     frontendTools: [{frontend_tools_str}],
-  environmentName: '{spec.get("environment_name", "ai-agents-env")}',
-  icon: {icon},
-  emoji: {emoji},
-  color: {color},
-  suggestions: {suggestions_str},
+    environmentName: '{spec.get("environment_name", "ai-agents-env")}',
+    icon: {icon},
+    emoji: {emoji},
+    color: {color},
+    suggestions: {suggestions_str},
     welcomeMessage: {_fmt_ts_literal(welcome_message)},
     welcomeNotebook: {_fmt_ts_literal(welcome_notebook)},
     welcomeDocument: {_fmt_ts_literal(welcome_document)},
-  sandboxVariant: {sandbox_variant_ts},
-  systemPrompt: {f"`{system_prompt}`" if system_prompt else "undefined"},
-  systemPromptCodemodeAddons: {f"`{system_prompt_codemode_addons}`" if system_prompt_codemode_addons else "undefined"},
-  goal: {goal_ts},
-  protocol: {protocol_ts},
-  uiExtension: {ui_ext_ts},
-  trigger: {_fmt_ts_literal(trigger_val)},
-  modelConfig: {_fmt_ts_literal(model_cfg)},
-  mcpServerTools: {_fmt_ts_literal(mcp_srv_tools)},
-  guardrails: {_fmt_ts_literal(guardrails_val)},
-  evals: {_fmt_ts_literal(evals_val)},
-  codemode: {_fmt_ts_literal(codemode_val)},
-  output: {_fmt_ts_literal(output_val)},
-  advanced: {_fmt_ts_literal(advanced_val)},
-  authorizationPolicy: {auth_policy_ts},
-  notifications: {_fmt_ts_literal(notifs)},
-  memory: {memory_ts},
+    sandboxVariant: {sandbox_variant_ts},
+    systemPrompt: {f"`{system_prompt}`" if system_prompt else "undefined"},
+    systemPromptCodemodeAddons: {f"`{system_prompt_codemode_addons}`" if system_prompt_codemode_addons else "undefined"},
+    goal: {goal_ts},
+    protocol: {protocol_ts},
+    uiExtension: {ui_ext_ts},
+    trigger: {_fmt_ts_literal(trigger_val)},
+    modelConfig: {_fmt_ts_literal(model_cfg)},
+    mcpServerTools: {_fmt_ts_literal(mcp_srv_tools)},
+    guardrails: {_fmt_ts_literal(guardrails_val)},
+    evals: {_fmt_ts_literal(evals_val)},
+    codemode: {_fmt_ts_literal(codemode_val)},
+    output: {_fmt_ts_literal(output_val)},
+    advanced: {_fmt_ts_literal(advanced_val)},
+    authorizationPolicy: {auth_policy_ts},
+    notifications: {_fmt_ts_literal(notifs)},
+    memory: {memory_ts},
+    preHooks: {_fmt_ts_literal(pre_hooks_val)},
+    postHooks: {_fmt_ts_literal(post_hooks_val)},
+    toolHooks: {_fmt_ts_literal(tool_hooks_val)},
+    parameters: {_fmt_ts_literal(parameters_val)},
+    subagents: {subagents_ts},
 }};
 
 """