Add artifact tags to all pipeline steps

- Add appropriate tags to all artifacts across pipeline steps
- Use add_tags() with artifact parameter matching Annotated names
- Tags help categorize and filter artifacts in the ZenML dashboard

Author: strickvl

deep_research/steps/approval_step.py

@@ -10,7 +10,7 @@
     summarize_research_progress,
 )
 from utils.pydantic_models import ApprovalDecision, ReflectionOutput
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 from zenml.client import Client
 
 logger = logging.getLogger(__name__)
@@ -300,6 +300,9 @@ def get_research_approval_step(
             }
         )
 
+        # Add tag to the approval decision artifact
+        add_tags(tags=["hitl"], artifact="approval_decision")
+
         return ApprovalDecision(
             approved=False,
             selected_queries=[],

deep_research/steps/collect_tracing_metadata_step.py

@@ -18,7 +18,7 @@
     get_trace_stats,
     get_traces_by_name,
 )
-from zenml import get_step_context, step
+from zenml import add_tags, get_step_context, step
 
 logger = logging.getLogger(__name__)
 
@@ -231,4 +231,10 @@ def collect_tracing_metadata_step(
                 )
             metadata.search_queries_count = search_queries_count
 
+    # Add tags to the artifacts
+    add_tags(tags=["state"], artifact="state")
+    add_tags(
+        tags=["exa", "tavily", "llm", "cost"], artifact="tracing_metadata"
+    )
+
     return state, metadata

deep_research/steps/cross_viewpoint_step.py

@@ -14,7 +14,7 @@
     ViewpointAnalysis,
     ViewpointTension,
 )
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 
 logger = logging.getLogger(__name__)
 
@@ -187,6 +187,9 @@ def cross_viewpoint_analysis_step(
             infer_artifact=True,
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "viewpoint"], artifact="analyzed_state")
+
         return state
 
     except Exception as e:
@@ -220,4 +223,7 @@ def cross_viewpoint_analysis_step(
             }
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "viewpoint"], artifact="analyzed_state")
+
         return state

deep_research/steps/execute_approved_searches_step.py

@@ -18,7 +18,7 @@
     SynthesizedInfo,
 )
 from utils.search_utils import search_and_extract_results
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 
 logger = logging.getLogger(__name__)
 
@@ -143,6 +143,9 @@ def execute_approved_searches_step(
             }
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "enhanced"], artifact="updated_state")
+
         return state
 
     # Execute approved searches
@@ -373,6 +376,9 @@ def execute_approved_searches_step(
             infer_artifact=True,
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "enhanced"], artifact="updated_state")
+
         return state
 
     except Exception as e:
@@ -413,4 +419,7 @@ def execute_approved_searches_step(
             }
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "enhanced"], artifact="updated_state")
+
         return state

deep_research/steps/generate_reflection_step.py

@@ -8,7 +8,7 @@
 )
 from utils.llm_utils import get_structured_llm_output
 from utils.pydantic_models import Prompt, ReflectionOutput, ResearchState
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 
 logger = logging.getLogger(__name__)
 
@@ -160,4 +160,7 @@ def generate_reflection_step(
         infer_artifact=True,
     )
 
+    # Add tags to the artifact
+    add_tags(tags=["reflection", "critique"], artifact="reflection_output")
+
     return reflection_output

deep_research/steps/initialize_prompts_step.py

@@ -10,7 +10,7 @@
 from materializers.prompt_materializer import PromptMaterializer
 from utils import prompts
 from utils.pydantic_models import Prompt
-from zenml import step
+from zenml import add_tags, step
 
 logger = logging.getLogger(__name__)
 
@@ -119,6 +119,27 @@ def initialize_prompts_step(
 
     logger.info(f"Loaded 9 individual prompts")
 
+    # add tags to all prompts
+    add_tags(tags=["prompt", "search"], artifact="search_query_prompt")
+    add_tags(
+        tags=["prompt", "generation"], artifact="query_decomposition_prompt"
+    )
+    add_tags(tags=["prompt", "generation"], artifact="synthesis_prompt")
+    add_tags(
+        tags=["prompt", "generation"], artifact="viewpoint_analysis_prompt"
+    )
+    add_tags(tags=["prompt", "generation"], artifact="reflection_prompt")
+    add_tags(
+        tags=["prompt", "generation"], artifact="additional_synthesis_prompt"
+    )
+    add_tags(
+        tags=["prompt", "generation"], artifact="conclusion_generation_prompt"
+    )
+    add_tags(
+        tags=["prompt", "generation"], artifact="executive_summary_prompt"
+    )
+    add_tags(tags=["prompt", "generation"], artifact="introduction_prompt")
+
     return (
         search_query_prompt,
         query_decomposition_prompt,

deep_research/steps/iterative_reflection_step.py

@@ -16,7 +16,7 @@
     SynthesizedInfo,
 )
 from utils.search_utils import search_and_extract_results
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 
 logger = logging.getLogger(__name__)
 
@@ -353,6 +353,9 @@ def iterative_reflection_step(
             infer_artifact=True,
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "reflected"], artifact="reflected_state")
+
         return state
 
     except Exception as e:
@@ -382,4 +385,7 @@ def iterative_reflection_step(
             }
         )
 
+        # Add tags to the artifact
+        add_tags(tags=["state", "reflected"], artifact="reflected_state")
+
         return state

deep_research/steps/merge_results_step.py

@@ -5,7 +5,7 @@
 
 from materializers.pydantic_materializer import ResearchStateMaterializer
 from utils.pydantic_models import ResearchState
-from zenml import get_step_context, log_metadata, step
+from zenml import add_tags, get_step_context, log_metadata, step
 from zenml.client import Client
 
 logger = logging.getLogger(__name__)
@@ -262,4 +262,7 @@ def merge_sub_question_results_step(
         infer_artifact=True,
     )
 
+    # Add tags to the artifact
+    add_tags(tags=["state", "merged"], artifact="merged_state")
+
     return merged_state

deep_research/steps/process_sub_question_step.py

@@ -17,7 +17,7 @@
     generate_search_query,
     search_and_extract_results,
 )
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 
 logger = logging.getLogger(__name__)
 
@@ -87,6 +87,8 @@ def process_sub_question_step(
         )
         # Return an empty state since there's no question to process
         sub_state.sub_questions = []
+        # Add tags to the artifact
+        add_tags(tags=["state", "sub-question"], artifact="output")
         return sub_state
 
     # Get the target sub-question
@@ -283,4 +285,7 @@ def process_sub_question_step(
         infer_artifact=True,
     )
 
+    # Add tags to the artifact
+    add_tags(tags=["state", "sub-question"], artifact="output")
+
     return sub_state

deep_research/steps/pydantic_final_report_step.py

@@ -23,7 +23,7 @@
     VIEWPOINT_ANALYSIS_TEMPLATE,
 )
 from utils.pydantic_models import Prompt, ResearchState
-from zenml import log_metadata, step
+from zenml import add_tags, log_metadata, step
 from zenml.types import HTMLString
 
 logger = logging.getLogger(__name__)
@@ -1092,6 +1092,9 @@ def pydantic_final_report_step(
         logger.info(
             "Final research report generated successfully with static template"
         )
+        # Add tags to the artifacts
+        add_tags(tags=["state", "final"], artifact="state")
+        add_tags(tags=["report", "html"], artifact="report_html")
         return state, HTMLString(html_content)
 
     # Otherwise use the LLM-generated approach
@@ -1187,6 +1190,9 @@ def pydantic_final_report_step(
         )
 
         logger.info("Final research report generated successfully")
+        # Add tags to the artifacts
+        add_tags(tags=["state", "final"], artifact="state")
+        add_tags(tags=["report", "html"], artifact="report_html")
         return state, HTMLString(html_content)
 
     except Exception as e:
@@ -1245,4 +1251,7 @@ def pydantic_final_report_step(
             infer_model=True,
         )
 
+        # Add tags to the artifacts
+        add_tags(tags=["state", "final"], artifact="state")
+        add_tags(tags=["report", "html"], artifact="report_html")
         return state, HTMLString(fallback_html)