From eb12a70a847bf8601082a83412b4f0d30e9ac472 Mon Sep 17 00:00:00 2001
From: David Kreiner <david.kreiner@harness.io>
Date: Wed, 15 Jan 2025 03:28:22 +0000
Subject: [PATCH] feat: [ML-542]: Add stage type to genai step context (#3266)

* feat: [ML-542]: Fix linting
* feat: [ML-542]: Add stage type to genai step context
---
 genai/genai.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/genai/genai.go b/genai/genai.go
index d8217bfeb..c64049769 100644
--- a/genai/genai.go
+++ b/genai/genai.go
@@ -55,7 +55,8 @@ func (PipelineContext) GetType() capabilities2.AIContextPayloadType {
 const AIContextPayloadTypeStepContext capabilities2.AIContextPayloadType = "other"
 
 type StepContext struct {
-	Yaml string `json:"step_yaml"`
+	Yaml      string `json:"step_yaml"`
+	StageType string `json:"stage_type"`
 }
 
 func (StepContext) GetName() string {