refactor to functor

tools/helm_test_gen.py

@@ -2,78 +2,86 @@ import re
 import yaml
 import os
 from typing import Optional
-from smolagents import tool
+from smolagents import Tool
 from jinja2 import Environment, FileSystemLoader
 from pydash import py_
 
 
-SEP = '''
-
----
-
-'''
-
-
-def get_yaml_value(d: dict[str], xpath: str, default: Optional[str] = None) -> str:
-    try:
-        if '.' in xpath:
-            h = xpath.split('.')[0]
-            t = '.'.join(xpath.split('.')[1:])
-            return get_yaml_value(d[h], t)
-        else:
-            return d[xpath]
-    except Exception as ex:
-        if default:
-            return default
-        else:
-            raise ex
-
-
-def bool_to_str(value: bool) -> str:
-    return str(value).lower()
-
-
-def get_object_by_kind(document: dict[str], kinds: tuple[str]) -> list[dict[str]]:
-    return [*py_.filter_(document, lambda x: x['kind'] in kinds)]
-
-
-def get_envs_by_ref(container: dict[str], ref: str, name: str) -> list[dict[str]]:
-    return [*py_.filter(container['env'], lambda x: x.get('valueFrom', {f'{ref}': {}}).get(f'{ref}', {'name': ''}).get('name') == name)]  # noqa: E501
-
-
-def get_binding_subjects(subjects: dict[str], name: str) -> list[dict[str]]:
-    return [*py_.filter(subjects, lambda x: x['name'] == name)]
-
-
-@tool
-def generate_helm_tests(document: str) -> str:
-  """Generate Helm test scenarios from a Kubernetes YAML document.
-
-  Args:
-      document (str): A YAML string containing one or more Kubernetes resource definitions.
-
-  Returns:
-      str:  A string containing the generated test scenarios, separated by a predefined
-      separator when multiple resources are processed.
-  """
-  doc = [*yaml.safe_load_all(document)]
-  result = ''
-  current_dir = os.path.dirname(os.path.abspath(__file__))
-  template_dir = os.path.join(current_dir, 'templates')
-
-  for item in doc:
-      name = item['metadata']['name']
-      kind = item['kind']
-      kindDesc = (re.sub(r'((?<=[a-z])[A-Z]|(?<!\A)[A-Z](?=[a-z]))', r' \1', kind)).lower()
-      env = Environment(loader=FileSystemLoader(template_dir))
-      env.globals['gyv'] = get_yaml_value
-      env.globals['bool_to_str'] = bool_to_str
-      env.globals['get_object_by_kind'] = get_object_by_kind
-      env.globals['get_env_by_ref'] = get_envs_by_ref
-      env.globals['get_binding_subjects'] = get_binding_subjects
-      template = env.get_template('scenario.jinja')
-      output = template.render(name=name, app='app', kind=kind, kindDesc=kindDesc, item=item, document=doc)
-      if result:
-        result += SEP
-      result += output
-  return result
+class GenerateHelmTest(Tool):
+  name = "generate_helm_tests"
+  description = "Performs a duckduckgo web search based on your query (think a Google search) then returns the top search results."
+  inputs = {'document': {'type': 'string', 'description': 'A YAML string containing one or more Kubernetes resource definitions.'}}
+  output_type = "string"
+
+  SEP = '''
+
+  ---
+
+  '''
+
+  
+  def get_yaml_value(self, d: dict[str], xpath: str, default: Optional[str] = None) -> str:
+      try:
+          if '.' in xpath:
+              h = xpath.split('.')[0]
+              t = '.'.join(xpath.split('.')[1:])
+              return self.get_yaml_value(d[h], t)
+          else:
+              return d[xpath]
+      except Exception as ex:
+          if default:
+              return default
+          else:
+              raise ex
+
+
+  @staticmethod
+  def bool_to_str(value: bool) -> str:
+      return str(value).lower()
+
+
+  @staticmethod
+  def get_object_by_kind(document: dict[str], kinds: tuple[str]) -> list[dict[str]]:
+      return [*py_.filter_(document, lambda x: x['kind'] in kinds)]
+
+  @staticmethod
+  def get_envs_by_ref(container: dict[str], ref: str, name: str) -> list[dict[str]]:
+      return [*py_.filter(container['env'], lambda x: x.get('valueFrom', {f'{ref}': {}}).get(f'{ref}', {'name': ''}).get('name') == name)]  # noqa: E501
+
+  @staticmethod
+  def get_binding_subjects(subjects: dict[str], name: str) -> list[dict[str]]:
+      return [*py_.filter(subjects, lambda x: x['name'] == name)]
+
+
+
+  def forward(self, document: str) -> str:
+    """Generate Helm test scenarios from a Kubernetes YAML document.
+
+    Args:
+        document (str): A YAML string containing one or more Kubernetes resource definitions.
+
+    Returns:
+        str:  A string containing the generated test scenarios, separated by a predefined
+        separator when multiple resources are processed.
+    """
+    doc = [*yaml.safe_load_all(document)]
+    result = ''
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    template_dir = os.path.join(current_dir, 'templates')
+
+    for item in doc:
+        name = item['metadata']['name']
+        kind = item['kind']
+        kindDesc = (re.sub(r'((?<=[a-z])[A-Z]|(?<!\A)[A-Z](?=[a-z]))', r' \1', kind)).lower()
+        env = Environment(loader=FileSystemLoader(template_dir))
+        env.globals['gyv'] = self.get_yaml_value
+        env.globals['bool_to_str'] = GenerateHelmTest.bool_to_str
+        env.globals['get_object_by_kind'] = GenerateHelmTest.get_object_by_kind
+        env.globals['get_env_by_ref'] = GenerateHelmTest.get_envs_by_ref
+        env.globals['get_binding_subjects'] = GenerateHelmTest.get_binding_subjects
+        template = env.get_template('scenario.jinja')
+        output = template.render(name=name, app='app', kind=kind, kindDesc=kindDesc, item=item, document=doc)
+        if result:
+          result += GenerateHelmTest.SEP
+        result += output
+    return result