Merge pull request #29 from oleks-dev/refactor_output_filters

output filters refactored, added regex_replace for sanitization

Author: oleks-dev

docs/reference/config/providers.md

@@ -67,11 +67,11 @@ Send prompts to a command via STDIN and read STDOUT (e.g., q chat, mlx_lm.genera
 qchat:
   provider_type: "stdin_consumer"
   mode: "flat"
-  call: "q"                      # executable file
+  call: "q"                           # executable file
   args: ["chat", "--no-interactive"]  # optional arguments
-  strip_output_prefix: "> "      # optional [str] - strip prefix text from step output
-  slice_output_start: 0          # optional [int] - slice step output text starting from N
-  slice_output_end: -1           # optional [int] - slice step output text ending at N (negative counts from the back)
+  filter:
+    strip: true             # strip leading and ending spaces
+    strip_prefix: "> "      # remove q char initial character and space
 ```
 
 #### MLX LM Generate  
@@ -85,5 +85,6 @@ qchat:
       - "/Users/guest/.cache/huggingface/hub/models--mlx-community--Mistral-7B-Instruct-v0.3-4bit/snapshots/a4b8f870474b0eb527f466a03fbc187830d271f5"
       - "--prompt"
       - "-"
-    output_regex: "^==========\\n((?:.|\\n)+)\\n\\=\\=\\=\\=\\=\\=\\=\\=\\=\\=(?:.|\\n)+$"
+    filter:
+      regex_extract: "^==========\\n((?:.|\\n)+)\\n\\=\\=\\=\\=\\=\\=\\=\\=\\=\\=(?:.|\\n)+$"
 ```

docs/reference/template/steps.md

@@ -21,22 +21,27 @@ steps:
 
 ##### Output text transformations  
 ```yaml
-    # strip spaces from beginning and end of the output
-    strip_output: true                 # optional [bool]
+    filter:
+      # strip spaces from beginning and end of the output
+      strip: true                 # optional [bool]
 
-    # strip prefix characters from the output
-    strip_output_prefix: "> "          # optional [str]
+      # strip prefix characters from the output
+      strip_prefix: "> "          # optional [str]
 
-    # slice output from N character
-    slice_output_start: 3              # optional [int]
+      # slice output from N character
+      slice_start: 3              # optional [int]
 
-    # slice output till N character 
-    # (use negative number for backwards count)
-    slice_output_end: -1               # optional [int]
+      # slice output till N character 
+      # (use negative number for backwards count)
+      slice_end: -1               # optional [int]
 
-    # filter output using regex 
-    # (take 1st group if groups are used otherwise take matching regex)
-    output_regex: ".*"                 # optional [str]
+      # filter output using regex 
+      # (take 1st group if groups are used otherwise take matching regex)
+      regex_extract: ".*"                 # optional [str]
+
+      # replace output text using regex patterns
+      regex_replace:          # optional [list(tuple(str,str))] - regex pattern, replace
+        - ["(?i)(\"password\"\s*:\s*\")[^\"]+(\")", "\\1*****\\2"]  # (ex. for json passwords sanitization)
 ```
 
 > **Note:** Output text transformation params applied one by one in order as they mentioned in the example, each next one uses result of the previous.  
@@ -48,10 +53,12 @@ steps:
     output_variable: "out_var"         # optional [str]
 
     # save output to file
-    output_file: "out.txt"             # optional [str]
-
-    # file mode overwrite or append
-    output_file_mode: "write"          # optional ["write"|"append"]
+    output_file: "out.txt"             # optional [str|dict] - write mode by default
+    # OR
+    output_file:
+      name: "out.txt"
+      # file mode write or append
+      mode: "write"          # optional ["write"|"append"]
 
     # print output to console during normal execution
     # (mostly for user reference when such information is needed)

prich/core/engine.py

@@ -95,37 +95,36 @@ def run_template(template_id, **kwargs):
                 raise click.ClickException(f"Step {step.type} type is not supported.")
 
             if is_verbose():
-                if step.extract_vars or step.output_regex or step.strip_output_prefix or step.slice_output_start or step.slice_output_end:
-                    console_print(f"[dim]Output: '{step_output}'[/dim]")
-            step.postprocess_extract_vars(output=step_output, variables=variables)
-            step_output = step.postprocess_output(output=step_output)
+                if step.extract_variables or step.filter:
+                    console_print(f"[dim]Output:\n{step_output}[/dim]")
+            step.postprocess_extract_vars(out=step_output, variables=variables)
+            step_output = step.postprocess_filter(out=step_output)
             if is_verbose():
-                if step.extract_vars:
-                    for spec in step.extract_vars:
-                        console_print(f"[dim]Inject \"{spec.regex}\" {f'({len(variables.get(spec.variable))} matches) ' if spec.multiple else ''}→ {spec.variable}: {f'{variables.get(spec.variable)}' if type(variables.get(spec.variable) == str) else variables.get(spec.variable)}[/dim]")
-                if step.strip_output is not None:
-                    console_print(f"[dim]Strip output spaces: {step.strip_output}[/dim]")
-                if step.strip_output_prefix:
-                    console_print(f"[dim]Strip output prefix: \"{step.strip_output_prefix}\"[/dim]")
-                if step.slice_output_start or step.slice_output_end:
-                    console_print(f"[dim]Slice output text{f' from {step.slice_output_start}' if step.slice_output_start else ''}{f' to {step.slice_output_end}' if step.slice_output_end else ''}[/dim]")
-                if step.output_regex:
-                    console_print(f"[dim]Apply regex: \"{step.output_regex}\"[/dim]")
+                if step.extract_variables:
+                    for spec in step.extract_variables:
+                        console_print(f"""[dim]Inject \"{spec.regex}\" {f'({len(variables.get(spec.variable))} matches) ' if spec.multiple else ''}→ {spec.variable}: {f'"{variables.get(spec.variable)}"' if type(variables.get(spec.variable)) == str else variables.get(spec.variable)}[/dim]""")
+                if step.filter:
+                    if step.filter.strip is not None:
+                        console_print(f"[dim]Strip output spaces: {step.filter.strip}[/dim]")
+                    if step.filter.strip_prefix:
+                        console_print(f"[dim]Strip output prefix: \"{step.filter.strip_prefix}\"[/dim]")
+                    if step.filter.slice_start or step.filter.slice_end:
+                        console_print(f"[dim]Slice output text{f' from {step.filter.slice_start}' if step.filter.slice_start else ''}{f' to {step.filter.slice_end}' if step.filter.slice_end else ''}[/dim]")
+                    if step.filter.regex_extract:
+                        console_print(f"[dim]Apply regex: \"{step.filter.regex_extract}\"[/dim]")
+                    if step.filter.regex_replace:
+                        replace_details = '\n                     '.join([f"\"{x}\" → \"{y}\"" for x,y in step.filter.regex_replace])
+                        console_print(f"[dim]Apply regex replace: {replace_details}[/dim]")
 
             # Store last output
             last_output = step_output
 
             if output_var:
                 variables[output_var] = step_output
             if step.output_file:
-                if step.output_file.startswith('.'):
-                    save_to_file = step.output_file.replace('.', str(get_cwd_dir()), 1)
-                elif step.output_file.startswith('~'):
-                    save_to_file = step.output_file.replace('~', str(get_home_dir()), 1)
-                else:
-                    save_to_file = step.output_file
+                save_to_file = step.output_file.name
                 try:
-                    write_mode = step.output_file_mode[:1] if step.output_file_mode else 'w'
+                    write_mode = step.output_file.mode[:1] if step.output_file.mode else 'w'
                     with open(save_to_file, write_mode) as step_output_file:
                         if is_verbose():
                             console_print(f"[dim]{'Save' if write_mode == 'w' else 'Append'} output to file: {save_to_file}[/dim]")

prich/core/steps/step_sent_to_llm.py

@@ -75,7 +75,7 @@ def send_to_llm(template: TemplateModel, step: LLMStep, provider: str, config: C
                 instructions=step.rendered_instructions,
                 input_=step.rendered_input
             )
-        step_output = selected_provider.postprocess_output(response)
+        step_output = selected_provider.postprocess_filter(response)
         if (is_verbose() or step.output_console) and not llm_provider.show_response and not is_quiet():
             console_print(step_output, markup=False)
     except Exception as e:

prich/models/config_providers.py

@@ -1,14 +1,22 @@
-from pydantic import Field, ConfigDict
+from pydantic import Field, ConfigDict, BaseModel
 from typing import Literal, Optional, List, Tuple
-from prich.models.output_shaping import BaseOutputShapingModel
+from prich.models.text_filter_model import TextFilterModel
 
 
-class BaseProviderModel(BaseOutputShapingModel):
+class BaseProviderModel(BaseModel):
     model_config = ConfigDict(extra='forbid')
     name: str | None = Field(default=None, exclude=True)  # will be injected
 
     mode: Optional[str] = None
 
+    # transforms
+    filter: Optional[TextFilterModel] = None
+
+    def postprocess_filter(self, out):
+        if self.filter:
+            out = self.filter.apply(out)
+        return out
+
     def model_post_init(self, __context):
         if self.name is None and __context and "__name" in __context:
             self.name = __context["__name"]

prich/models/output_shaping.py

@@ -1,11 +1,10 @@
-from pathlib import Path
-
+import re
 import click
-from pydantic import BaseModel, Field, model_validator, ConfigDict
+from pydantic import BaseModel, Field, model_validator, ConfigDict, field_validator
 from typing import List, Optional, Literal, Annotated, Union
-from prich.models.output_shaping import BaseOutputShapingModel
-from prich.constants import RESERVED_RUN_TEMPLATE_CLI_OPTIONS
+from prich.models.text_filter_model import TextFilterModel
 from prich.models.file_scope import FileScope
+from prich.constants import RESERVED_RUN_TEMPLATE_CLI_OPTIONS
 from prich.core.utils import is_valid_variable_name, is_cli_option_name, get_prich_dir
 from prich.version import TEMPLATE_SCHEMA_VERSION
 
@@ -45,46 +44,70 @@ class ExtractVarModel(BaseModel):
     variable: str
     multiple: Optional[bool] = False  # default: single match
 
+    def extract(self, text: str) -> str | list:
+        pattern = re.compile(self.regex)
+        if self.multiple:
+            matches = pattern.findall(text)
+            if matches:
+                # if regex has groups, findall returns tuples
+                values = [m if isinstance(m, str) else m[0] for m in matches]
+                return values
+            return []
+        else:
+            m = pattern.search(text)
+            if m:
+                return m.group(1) if m.groups() else m.group(0)
+            return ""
+
+
+class OutputFileModel(BaseModel):
+    name: Optional[str | None] = None
+    mode: Optional[Literal["write", "append"] | None] = None
 
-class BaseStepModel(BaseOutputShapingModel):
+
+class BaseStepModel(BaseModel):
     model_config = ConfigDict(extra='forbid')
 
     name: str
 
     # regex transforms
-    extract_vars: Optional[list[ExtractVarModel]] = None  # enrichment
+    extract_variables: Optional[list[ExtractVarModel]] = None  # enrichment
+
+    # output transforms
+    filter: Optional[TextFilterModel] = None
 
     # persistence
     output_variable: Optional[str | None] = None
-    output_file: Optional[str | None] = None
-    output_file_mode: Optional[Literal["write", "append"]] = None
+    output_file: Optional[str | OutputFileModel | None] = None
     output_console: Optional[bool | None] = None
 
     # execution control
     when: Optional[str | None] = None
     validate_: Optional[ValidateStepOutput | list[ValidateStepOutput]] = Field(alias="validate", default=None)
 
-    def postprocess_extract_vars(self, output: str, variables: dict):
-        import re
-
+    # normalize output_file
+    @field_validator("output_file")
+    def normalize_output_file(cls, v):
+        if v is None:
+            return None
+        if isinstance(v, OutputFileModel):
+            return v
+        if isinstance(v, str):
+            return OutputFileModel(name=v)
+        if isinstance(v, dict):
+            return OutputFileModel(name=v.get("name"), mode=v.get("mode", None))
+        raise ValueError("Invalid format for variable field in TransformStep")
+
+    def postprocess_extract_vars(self, out: str, variables: dict):
         # extract side variables
-        if self.extract_vars:
-            for spec in self.extract_vars:
-                pattern = re.compile(spec.regex)
-                if spec.multiple:
-                    matches = pattern.findall(output)
-                    if matches:
-                        # if regex has groups, findall returns tuples
-                        values = [m if isinstance(m, str) else m[0] for m in matches]
-                        variables[spec.variable] = values
-                    else:
-                        variables[spec.variable] = []
-                else:
-                    m = pattern.search(output)
-                    if m:
-                        variables[spec.variable] = m.group(1) if m.groups() else m.group(0)
-                    else:
-                        variables[spec.variable] = ""
+        if self.extract_variables:
+            for spec in self.extract_variables:
+                variables[spec.variable] = spec.extract(out)
+
+    def postprocess_filter(self, out):
+        if self.filter:
+            out = self.filter.apply(out)
+        return out
 
 
 class LLMStep(BaseStepModel):

prich/models/template.py

@@ -0,0 +1,33 @@
+import re
+from typing import Optional, List, Tuple
+from pydantic import BaseModel
+
+class TextFilterModel(BaseModel):
+    strip: Optional[bool] = True
+    strip_prefix: Optional[str] = None
+    slice_start: Optional[int] = None
+    slice_end: Optional[int] = None
+    regex_extract: Optional[str] = None
+    regex_replace: Optional[List[Tuple[str, str]]] = None  # [(pattern, replacement), ...]
+
+    def apply(self, text: str) -> str:
+        out = text
+
+        if self.strip:
+            out = out.strip()
+
+        if self.strip_prefix and out.startswith(self.strip_prefix):
+            out = out[len(self.strip_prefix):]
+
+        if self.slice_start is not None or self.slice_end is not None:
+            out = out[self.slice_start:self.slice_end]
+
+        if self.regex_extract:
+            m = re.search(self.regex_extract, out)
+            out = m.group(1) if (m and m.groups()) else (m.group(0) if m else "")
+
+        if self.regex_replace:
+            for pattern, repl in self.regex_replace:
+                out = re.sub(pattern, repl, out)
+
+        return out

prich/models/text_filter_model.py

@@ -27,11 +27,14 @@ Each step defines an action for the template pipeline workflow executed in order
 **Step can have next key parameters**:
 * `name`: name of the step (should be unique for each step in the template)
 * `output_variable`: save output of the execution into a variable for the following usage
-* `output_file`: save output of the execution into a file
-* `output_file_mode`: `write`/`append`
-* `strip_output_prefix`: strip prefix string from the step output
-* `slice_output_start`: slice step output from character number
-* `slice_output_end`: slice step output to character number
+* `output_file`: save output of the execution into a file (default would be `write` mode)
+* `output_file`: 
+    `name`: save output of the execution into a file
+    `mode`: `write` or `append`
+* `filter`:
+    `strip_prefix`: strip prefix string from the step output
+    `slice_start`: slice step output from character number
+    `slice_end`: slice step output to character number
 * `when`: execute step only when true - simple jinja evaluation like `working_vs_last_commit or (not working_vs_last_commit and not working_vs_remote and not committed_vs_remote and not remote_vs_local)` or `not remote_vs_local`, etc.
 * `validate`: validate step execution (see `step validate`)
 * plus additional keys based on the step type