fix: [ML-505]: changes to support update-step API for Harness Intelligence (#3156)

* Fix some lint issues
* fix more lint issues
* Fix linting
* fix: [ML-505]: add step context
* fix: [ML-505]: fix interface for update step
* fix: [ML-505]: interface for update step
* fix: [ML-505]: interface for update step
* fix: [ML-505]: structs for update step

app/api/controller/aiagent/generate_pipeline.go

@@ -27,8 +27,7 @@ func (c *Controller) GeneratePipeline(
 	in *controllertypes.GeneratePipelineInput,
 ) (*controllertypes.GeneratePipelineOutput, error) {
 	generateRequest := &aitypes.PipelineGenerateRequest{
-		Prompt:  in.Prompt,
-		RepoRef: in.RepoRef,
+		Prompt: in.Prompt,
 	}
 
 	output, err := c.intelligence.GeneratePipeline(ctx, generateRequest)
@@ -36,7 +35,6 @@ func (c *Controller) GeneratePipeline(
 		return nil, fmt.Errorf("generate pipeline: %w", err)
 	}
 	return &controllertypes.GeneratePipelineOutput{
-		Status: "SUCCESS",
 		Data: controllertypes.PipelineData{
 			YamlPipeline: output.YAML,
 		},

app/api/controller/aiagent/generate_pipeline_step.go

@@ -27,8 +27,7 @@ func (c *Controller) GeneratePipelineStep(
 	in *controllertypes.GeneratePipelineStepInput,
 ) (*controllertypes.GeneratePipelineStepOutput, error) {
 	generateRequest := &aitypes.PipelineStepGenerateRequest{
-		Prompt:  in.Prompt,
-		RepoRef: in.RepoRef,
+		Prompt: in.Prompt,
 	}
 
 	output, err := c.intelligence.GeneratePipelineStep(ctx, generateRequest)
@@ -36,9 +35,8 @@ func (c *Controller) GeneratePipelineStep(
 		return nil, fmt.Errorf("generate pipeline: %w", err)
 	}
 	return &controllertypes.GeneratePipelineStepOutput{
-		Status: "SUCCESS",
 		Data: controllertypes.PipelineStepData{
-			StepYaml: output.YAML,
+			Yaml: output.Yaml,
 		},
 	}, nil
 }

app/api/controller/aiagent/types/pipeline.go

@@ -14,22 +14,22 @@
 
 package types
 
-type GeneratePipelineInput struct {
-	Prompt  string `json:"prompt"`
-	RepoRef string `json:"repo_ref"`
-}
-
 type PipelineData struct {
 	YamlPipeline string `json:"yaml_pipeline"`
 }
 
-type GeneratePipelineOutput struct {
-	Status string       `json:"status"`
-	Data   PipelineData `json:"data"`
+// create.
+type GeneratePipelineInput struct {
+	Prompt string `json:"prompt"`
 }
 
+type GeneratePipelineOutput struct {
+	Error string       `json:"error"`
+	Data  PipelineData `json:"data"`
+}
+
+// suggest.
 type SuggestPipelineInput struct {
-	RepoRef  string `json:"repo_ref"`
 	Pipeline string `json:"pipeline"`
 }
 
@@ -37,13 +37,13 @@ type SuggestPipelineOutput struct {
 	Suggestions []Suggestion
 }
 
-type UpdatePipelineOutput struct {
-	Status string       `json:"status"`
-	Data   PipelineData `json:"data"`
-}
-
+// update.
 type UpdatePipelineInput struct {
 	Prompt   string `json:"prompt"`
-	RepoRef  string `json:"repo_ref"`
 	Pipeline string `json:"pipeline"`
 }
+
+type UpdatePipelineOutput struct {
+	Error string       `json:"error"`
+	Data  PipelineData `json:"data"`
+}

app/api/controller/aiagent/types/stage.go

@@ -24,6 +24,6 @@ type PipelineStageData struct {
 }
 
 type GeneratePipelineStageOutput struct {
-	Status string            `json:"status"`
-	Data   PipelineStageData `json:"data"`
+	Error string            `json:"error"`
+	Data  PipelineStageData `json:"data"`
 }

app/api/controller/aiagent/types/step.go

@@ -14,16 +14,27 @@
 
 package types
 
-type GeneratePipelineStepInput struct {
-	Prompt  string `json:"prompt"`
-	RepoRef string `json:"repo_ref"`
+type PipelineStepData struct {
+	Yaml string `json:"yaml"`
 }
 
-type PipelineStepData struct {
-	StepYaml string `json:"yaml_step"`
+// create.
+type GeneratePipelineStepInput struct {
+	Prompt string `json:"prompt"`
 }
 
 type GeneratePipelineStepOutput struct {
-	Status string           `json:"status"`
+	Error string           `json:"error"`
+	Data  PipelineStepData `json:"data"`
+}
+
+// update.
+type UpdatePipelineStepInput struct {
+	Prompt string           `json:"prompt"`
 	Data   PipelineStepData `json:"data"`
 }
+
+type UpdatePipelineStepOutput struct {
+	Error string           `json:"error"`
+	Data  PipelineStepData `json:"data"`
+}

app/api/controller/aiagent/update_pipeline_step.go

@@ -0,0 +1,43 @@
+// Copyright 2023 Harness, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package aiagent
+
+import (
+	"context"
+	"fmt"
+
+	controllertypes "github.com/harness/gitness/app/api/controller/aiagent/types"
+	aitypes "github.com/harness/gitness/types/aigenerate"
+)
+
+func (c *Controller) UpdatePipelineStep(
+	ctx context.Context,
+	in *controllertypes.UpdatePipelineStepInput,
+) (*controllertypes.UpdatePipelineStepOutput, error) {
+	generateRequest := &aitypes.PipelineStepUpdateRequest{
+		Prompt: in.Prompt,
+		Step:   in.Data.Yaml,
+	}
+
+	output, err := c.intelligence.UpdatePipelineStep(ctx, generateRequest)
+	if err != nil {
+		return nil, fmt.Errorf("update pipeline: %w", err)
+	}
+	return &controllertypes.UpdatePipelineStepOutput{
+		Data: controllertypes.PipelineStepData{
+			Yaml: output.Yaml,
+		},
+	}, nil
+}

app/services/aiagent/intelligence.go

@@ -27,15 +27,36 @@ var _ Intelligence = GitnessIntelligence{}
 
 // This interface serves are the single interface to provide AI use cases.
 type Intelligence interface {
-	GeneratePipeline(
-		ctx context.Context,
-		req *aitypes.PipelineGenerateRequest) (*aitypes.PipelineGenerateResponse, error)
-	GeneratePipelineStage(
-		ctx context.Context,
-		req *aitypes.PipelineStageGenerateRequest) (*aitypes.PipelineStageGenerateResponse, error)
+	StepIntelligence
+	StageIntelligence
+	PipelineIntelligence
+}
+
+type StepIntelligence interface {
 	GeneratePipelineStep(
 		ctx context.Context,
 		req *aitypes.PipelineStepGenerateRequest) (*aitypes.PipelineStepGenerateResponse, error)
+	UpdatePipelineStep(
+		ctx context.Context,
+		req *aitypes.PipelineStepUpdateRequest) (*aitypes.PipelineStepUpdateResponse, error)
+}
+
+type StageIntelligence interface {
+	GeneratePipelineStage(
+		ctx context.Context,
+		req *aitypes.PipelineStageGenerateRequest) (*aitypes.PipelineStageGenerateResponse, error)
+	UpdatePipelineStage(
+		ctx context.Context,
+		req *aitypes.PipelineStageUpdateRequest) (*aitypes.PipelineStageUpdateResponse, error)
+}
+
+type PipelineIntelligence interface {
+	GeneratePipeline(
+		ctx context.Context,
+		req *aitypes.PipelineGenerateRequest) (*aitypes.PipelineGenerateResponse, error)
+	UpdatePipeline(
+		ctx context.Context,
+		req *aitypes.PipelineUpdateResponse) (*aitypes.PipelineUpdateResponse, error)
 }
 
 type GitnessIntelligence struct {
@@ -44,6 +65,27 @@ type GitnessIntelligence struct {
 	cc         *capabilitiesctrl.Controller
 }
 
+// UpdatePipeline implements Intelligence.
+func (h GitnessIntelligence) UpdatePipeline(
+	_ context.Context,
+	_ *aitypes.PipelineUpdateResponse) (*aitypes.PipelineUpdateResponse, error) {
+	panic("unimplemented")
+}
+
+// UpdatePipelineStage implements Intelligence.
+func (h GitnessIntelligence) UpdatePipelineStage(
+	_ context.Context,
+	_ *aitypes.PipelineStageUpdateRequest) (*aitypes.PipelineStageUpdateResponse, error) {
+	panic("unimplemented")
+}
+
+// UpdatePipelineStep implements Intelligence.
+func (h GitnessIntelligence) UpdatePipelineStep(
+	_ context.Context,
+	_ *aitypes.PipelineStepUpdateRequest) (*aitypes.PipelineStepUpdateResponse, error) {
+	panic("unimplemented")
+}
+
 // GeneratePipeline implements Intelligence.
 func (h GitnessIntelligence) GeneratePipeline(
 	_ context.Context,

genai/genai.go

@@ -15,16 +15,6 @@
 package genai
 
 import (
-	"bytes"
-	"context"
-	"encoding/json"
-	"io"
-	"net/http"
-	"os"
-
-	capabilitiesctrl "github.com/harness/gitness/app/api/controller/capabilities"
-	capabilitieshandler "github.com/harness/gitness/app/api/handler/capabilities"
-	"github.com/harness/gitness/app/services/capabilities"
 	capabilities2 "github.com/harness/gitness/types/capabilities"
 )
 
@@ -44,6 +34,10 @@ func GenerateAIContext(payloads ...capabilities2.AIContextPayload) []capabilitie
 	return out
 }
 
+var _ capabilities2.AIContextPayload = (*PipelineContext)(nil)
+var _ capabilities2.AIContextPayload = (*StepContext)(nil)
+var _ capabilities2.AIContextPayload = (*RepoRef)(nil)
+
 type PipelineContext struct {
 	Yaml string `json:"pipeline_yaml"`
 }
@@ -58,6 +52,20 @@ func (PipelineContext) GetType() capabilities2.AIContextPayloadType {
 	return AIContextPayloadTypePipelineContext
 }
 
+const AIContextPayloadTypeStepContext capabilities2.AIContextPayloadType = "other"
+
+type StepContext struct {
+	Yaml string `json:"step_yaml"`
+}
+
+func (StepContext) GetName() string {
+	return "step_context"
+}
+
+func (StepContext) GetType() capabilities2.AIContextPayloadType {
+	return AIContextPayloadTypeStepContext
+}
+
 type RepoRef struct {
 	Ref string `json:"ref"`
 }
@@ -71,44 +79,3 @@ const AIContextPayloadTypeRepoRef capabilities2.AIContextPayloadType = "other"
 func (RepoRef) GetType() capabilities2.AIContextPayloadType {
 	return AIContextPayloadTypeRepoRef
 }
-
-type ChatRequest struct {
-	Prompt          string                              `json:"prompt"`
-	ConversationID  string                              `json:"conversation_id"`
-	ConversationRaw string                              `json:"conversation_raw"`
-	Context         []capabilities2.AIContext           `json:"context"`
-	Capabilities    []capabilities2.CapabilityReference `json:"capabilities"`
-}
-
-func CallAIFoundation(ctx context.Context, cr *capabilities.Registry,
-	req *ChatRequest) (*capabilitiesctrl.RunCapabilitiesRequest, error) {
-	url := "http://localhost:8000/chat/platform"
-
-	jsonData, err := json.Marshal(req)
-	if err != nil {
-		return nil, err
-	}
-	newReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData))
-	if err != nil {
-		return nil, err
-	}
-	newReq.Header.Add("Authorization", "Bearer "+os.Getenv(AIFoundationServiceToken))
-
-	client := http.DefaultClient
-	resp, err := client.Do(newReq)
-	if err != nil {
-		return nil, err
-	}
-	defer resp.Body.Close()
-
-	data, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return nil, err
-	}
-	in, err := capabilitieshandler.UnmarshalRunCapabilitiesRequest(cr, data)
-	if err != nil {
-		return nil, err
-	}
-
-	return in, nil
-}

types/aigenerate/step.go

@@ -15,20 +15,18 @@
 package types
 
 type PipelineStepGenerateRequest struct {
-	Prompt  string
-	RepoRef string
+	Prompt string
 }
 
 type PipelineStepGenerateResponse struct {
-	YAML string
+	Yaml string
 }
 
 type PipelineStepUpdateRequest struct {
-	Prompt  string
-	RepoRef string
-	Step    string
+	Prompt string
+	Step   string
 }
 
 type PipelineStepUpdateResponse struct {
-	YAML string
+	Yaml string
 }