diff --git a/athina/helpers/step_helper.py b/athina/helpers/step_helper.py new file mode 100644 index 0000000..f478d85 --- /dev/null +++ b/athina/helpers/step_helper.py @@ -0,0 +1,10 @@ +import json + +class StepHelper: + + @staticmethod + def prepare_input_data(data): + return { + key: json.dumps(value) if isinstance(value, (list, dict)) else value + for key, value in data.items() + } \ No newline at end of file diff --git a/athina/steps/loop.py b/athina/steps/loop.py index 869c903..f25dce2 100644 --- a/athina/steps/loop.py +++ b/athina/steps/loop.py @@ -1,8 +1,11 @@ import asyncio +import json from typing import Dict, List, Any, Optional from athina.steps.base import Step from concurrent.futures import ThreadPoolExecutor - +from jinja2 import Environment +from athina.helpers.jinja_helper import PreserveUndefined +from athina.helpers.step_helper import StepHelper class Loop(Step): loop_type: str @@ -50,7 +53,16 @@ async def _execute_loop(self, inputs: Dict) -> Dict: results = [] if self.loop_type == "map": - items = inputs.get(self.loop_input, []) + env = Environment( + variable_start_string="{{", + variable_end_string="}}", + undefined=PreserveUndefined, + ) + + loop_input_template = env.from_string(self.loop_input) + prepared_input_data = StepHelper.prepare_input_data(inputs) + loop_input = loop_input_template.render(**prepared_input_data) + items = json.loads(loop_input, strict=False) if loop_input else None if not isinstance(items, list): return {"status": "error", "data": "Input not of type list", "metadata": {}} diff --git a/pyproject.toml b/pyproject.toml index 16a4f93..b34c689 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "athina" -version = "1.7.13" +version = "1.7.14" description = "Python SDK to configure and run evaluations for your LLM-based application" authors = ["Shiv Sakhuja ", "Akshat Gupta ", "Vivek Aditya ", "Akhil Bisht "] readme = "README.md"