homelab-frontend/internal/proxy/workflows.go

// Package proxy provides request routing and forwarding.
package proxy

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// WorkflowRequest represents a workflow execution request
type WorkflowRequest struct {
	// Workflow ID or name
	Workflow string `json:"workflow"`
	
	// Input parameters for the workflow
	Input map[string]interface{} `json:"input"`
	
	// Optional: timeout in seconds
	Timeout int `json:"timeout,omitempty"`
	
	// Optional: wait for result (default: true)
	Wait *bool `json:"wait,omitempty"`
}

// WorkflowResponse represents the response from workflow execution
type WorkflowResponse struct {
	// Workflow execution ID
	ID string `json:"id"`
	
	// Workflow name
	Workflow string `json:"workflow"`
	
	// Execution status: pending, running, completed, failed
	Status string `json:"status"`
	
	// Output of the workflow
	Output interface{} `json:"output,omitempty"`
	
	// Error message if workflow failed
	Error string `json:"error,omitempty"`
	
	// Timestamp when workflow was created
	CreatedAt time.Time `json:"created_at"`
	
	// Timestamp when workflow completed
	CompletedAt *time.Time `json:"completed_at,omitempty"`
}

// PredefinedWorkflow defines a workflow template that combines multiple API calls
type PredefinedWorkflow struct {
	Name        string
	Description string
	Handler     func(*http.Request, *Handler, map[string]interface{}) (interface{}, error)
}

// handleWorkflow handles the /workflows endpoint
// It accepts workflow definitions and orchestrates API calls
func (h *Handler) handleWorkflow(w http.ResponseWriter, r *http.Request) {
	// Only POST is supported
	if r.Method != "POST" {
		w.Header().Set("Content-Type", "application/problem+json")
		w.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(w, `{"type":"https://api.example.com/problems/method-not-allowed","title":"Method Not Allowed","status":405,"detail":"Only POST is supported for /workflows"}`)
		return
	}

	// Parse request body
	var workflowReq WorkflowRequest
	if err := json.NewDecoder(r.Body).Decode(&workflowReq); err != nil {
		writeProblemDetail(w, http.StatusBadRequest, "https://api.example.com/problems/invalid-workflow-request", "Invalid Workflow Request", "Failed to parse workflow request: "+err.Error(), nil)
		return
	}

	// Validate workflow name
	if workflowReq.Workflow == "" {
		writeProblemDetail(w, http.StatusBadRequest, "https://api.example.com/problems/missing-workflow", "Missing Workflow", "The 'workflow' field is required", nil)
		return
	}

	// Get predefined workflow
	workflow, ok := h.getWorkflow(workflowReq.Workflow)
	if !ok {
		availableWorkflows := h.getAvailableWorkflows()
		writeProblemDetail(w, http.StatusBadRequest, "https://api.example.com/problems/unknown-workflow", "Unknown Workflow", fmt.Sprintf("Workflow %q is not available", workflowReq.Workflow), availableWorkflows)
		return
	}

	// Default wait to true
	wait := true
	if workflowReq.Wait != nil {
		wait = *workflowReq.Wait
	}

	// Set default timeout if not provided
	timeout := time.Duration(30) * time.Second
	if workflowReq.Timeout > 0 {
		timeout = time.Duration(workflowReq.Timeout) * time.Second
	}

	// Create a context with timeout for workflow execution
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// Execute workflow
	output, err := workflow.Handler(r.WithContext(ctx), h, workflowReq.Input)

	// Build response
	workflowResp := WorkflowResponse{
		ID:        generateWorkflowID(),
		Workflow:  workflowReq.Workflow,
		CreatedAt: time.Now(),
	}

	if err != nil {
		workflowResp.Status = "failed"
		workflowResp.Error = err.Error()
	} else {
		if wait {
			workflowResp.Status = "completed"
			workflowResp.Output = output
			now := time.Now()
			workflowResp.CompletedAt = &now
		} else {
			workflowResp.Status = "pending"
		}
	}

	// Write response
	w.Header().Set("Content-Type", "application/json")
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	} else {
		w.WriteHeader(http.StatusOK)
	}
	json.NewEncoder(w).Encode(workflowResp)
}

// getWorkflow returns a predefined workflow by name
func (h *Handler) getWorkflow(name string) (*PredefinedWorkflow, bool) {
	workflows := h.getPredefinedWorkflows()
	for _, wf := range workflows {
		if wf.Name == name {
			return &wf, true
		}
	}
	return nil, false
}

// getPredefinedWorkflows returns all available workflows
func (h *Handler) getPredefinedWorkflows() []PredefinedWorkflow {
	return []PredefinedWorkflow{
		{
			Name:        "chat-and-embed",
			Description: "Chat with a model and then embed the response",
			Handler:     h.chatAndEmbedWorkflow,
		},
		{
			Name:        "multi-model-chat",
			Description: "Chat with multiple models sequentially",
			Handler:     h.multiModelChatWorkflow,
		},
		{
			Name:        "rag-pipeline",
			Description: "RAG pipeline: embed query, rerank, then chat with context",
			Handler:     h.ragPipelineWorkflow,
		},
		{
			Name:        "batch-embeddings",
			Description: "Generate embeddings for multiple texts",
			Handler:     h.batchEmbeddingsWorkflow,
		},
	}
}

// getAvailableWorkflows returns a list of available workflow names
func (h *Handler) getAvailableWorkflows() []string {
	workflows := h.getPredefinedWorkflows()
	names := make([]string, len(workflows))
	for i, wf := range workflows {
		names[i] = wf.Name
	}
	return names
}

// Workflow implementations

// chatAndEmbedWorkflow: Chat with a model, then embed the response
func (h *Handler) chatAndEmbedWorkflow(r *http.Request, handler *Handler, input map[string]interface{}) (interface{}, error) {
	model, ok := input["model"].(string)
	if !ok || model == "" {
		return nil, fmt.Errorf("missing required parameter: model")
	}

	embedModel, ok := input["embed_model"].(string)
	if !ok {
		embedModel = "nomic-ai/nomic-embed-text-v2-moe"
	}

	messages, ok := input["messages"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("missing required parameter: messages")
	}

	// Step 1: Chat
	chatReq := map[string]interface{}{
		"model":    model,
		"messages": messages,
	}

	chatBody, _ := json.Marshal(chatReq)
	chatHTTPReq, _ := http.NewRequest("POST", "/v1/chat/completions", io.NopCloser(bytes.NewReader(chatBody)))
	chatHTTPReq.Header.Set("Content-Type", "application/json")

	// Create a response writer to capture the chat response
	chatResp := &responseCapture{}
	handler.ServeHTTP(chatResp, chatHTTPReq)

	var chatResult map[string]interface{}
	if err := json.Unmarshal(chatResp.body.Bytes(), &chatResult); err != nil {
		return nil, fmt.Errorf("failed to parse chat response: %v", err)
	}

	// Extract message content
	var messageContent string
	if choices, ok := chatResult["choices"].([]interface{}); ok && len(choices) > 0 {
		if choice, ok := choices[0].(map[string]interface{}); ok {
			if message, ok := choice["message"].(map[string]interface{}); ok {
				if content, ok := message["content"].(string); ok {
					messageContent = content
				}
			}
		}
	}

	// Step 2: Embed the response
	embedReq := map[string]interface{}{
		"model": embedModel,
		"input": messageContent,
	}

	embedBody, _ := json.Marshal(embedReq)
	embedHTTPReq, _ := http.NewRequest("POST", "/v1/embeddings", io.NopCloser(bytes.NewReader(embedBody)))
	embedHTTPReq.Header.Set("Content-Type", "application/json")

	embedResp := &responseCapture{}
	handler.ServeHTTP(embedResp, embedHTTPReq)

	var embedResult map[string]interface{}
	if err := json.Unmarshal(embedResp.body.Bytes(), &embedResult); err != nil {
		return nil, fmt.Errorf("failed to parse embedding response: %v", err)
	}

	return map[string]interface{}{
		"chat_response":       chatResult,
		"embedding_response":  embedResult,
	}, nil
}

// multiModelChatWorkflow: Chat with multiple models sequentially
func (h *Handler) multiModelChatWorkflow(r *http.Request, handler *Handler, input map[string]interface{}) (interface{}, error) {
	models, ok := input["models"].([]interface{})
	if !ok || len(models) == 0 {
		return nil, fmt.Errorf("missing required parameter: models (array)")
	}

	messages, ok := input["messages"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("missing required parameter: messages")
	}

	results := make([]map[string]interface{}, 0)

	for _, modelInterface := range models {
		model, ok := modelInterface.(string)
		if !ok {
			continue
		}

		chatReq := map[string]interface{}{
			"model":    model,
			"messages": messages,
		}

		chatBody, _ := json.Marshal(chatReq)
		chatHTTPReq, _ := http.NewRequest("POST", "/v1/chat/completions", io.NopCloser(bytes.NewReader(chatBody)))
		chatHTTPReq.Header.Set("Content-Type", "application/json")

		chatResp := &responseCapture{}
		handler.ServeHTTP(chatResp, chatHTTPReq)

		var chatResult map[string]interface{}
		if err := json.Unmarshal(chatResp.body.Bytes(), &chatResult); err != nil {
			results = append(results, map[string]interface{}{
				"model": model,
				"error": err.Error(),
			})
			continue
		}

		results = append(results, map[string]interface{}{
			"model":  model,
			"result": chatResult,
		})
	}

	return results, nil
}

// ragPipelineWorkflow: RAG pipeline - embed query, rerank, chat with context
func (h *Handler) ragPipelineWorkflow(r *http.Request, handler *Handler, input map[string]interface{}) (interface{}, error) {
	query, ok := input["query"].(string)
	if !ok || query == "" {
		return nil, fmt.Errorf("missing required parameter: query")
	}

	documents, ok := input["documents"].([]interface{})
	if !ok {
		return nil, fmt.Errorf("missing required parameter: documents")
	}

	model, ok := input["model"].(string)
	if !ok {
		model = "reasoning"
	}

	rerankModel, ok := input["rerank_model"].(string)
	if !ok {
		rerankModel = "BAAI/bge-reranker-base"
	}

	topK := 3
	if tk, ok := input["top_k"].(float64); ok {
		topK = int(tk)
	}

	// Step 1: Rerank documents based on query
	rerankReq := map[string]interface{}{
		"model": rerankModel,
		"query": query,
		"texts": documents,
		"top_k": topK,
	}

	rerankBody, _ := json.Marshal(rerankReq)
	rerankHTTPReq, _ := http.NewRequest("POST", "/v1/rerank", io.NopCloser(bytes.NewReader(rerankBody)))
	rerankHTTPReq.Header.Set("Content-Type", "application/json")

	rerankResp := &responseCapture{}
	handler.ServeHTTP(rerankResp, rerankHTTPReq)

	var rerankResult map[string]interface{}
	if err := json.Unmarshal(rerankResp.body.Bytes(), &rerankResult); err != nil {
		return nil, fmt.Errorf("failed to parse rerank response: %v", err)
	}

	// Extract top documents
	var topDocs []string
	if results, ok := rerankResult["results"].([]interface{}); ok {
		for i, resultInterface := range results {
			if i >= topK {
				break
			}
			if result, ok := resultInterface.(map[string]interface{}); ok {
				if text, ok := result["text"].(string); ok {
					topDocs = append(topDocs, text)
				}
			}
		}
	}

	// Step 2: Chat with context
	context := fmt.Sprintf("Context from documents:\n%v\n\nQuery: %s", topDocs, query)

	chatReq := map[string]interface{}{
		"model": model,
		"messages": []interface{}{
			map[string]interface{}{
				"role":    "user",
				"content": context,
			},
		},
	}

	chatBody, _ := json.Marshal(chatReq)
	chatHTTPReq, _ := http.NewRequest("POST", "/v1/chat/completions", io.NopCloser(bytes.NewReader(chatBody)))
	chatHTTPReq.Header.Set("Content-Type", "application/json")

	chatResp := &responseCapture{}
	handler.ServeHTTP(chatResp, chatHTTPReq)

	var chatResult map[string]interface{}
	if err := json.Unmarshal(chatResp.body.Bytes(), &chatResult); err != nil {
		return nil, fmt.Errorf("failed to parse chat response: %v", err)
	}

	return map[string]interface{}{
		"reranked_documents": topDocs,
		"chat_response":      chatResult,
	}, nil
}

// batchEmbeddingsWorkflow: Generate embeddings for multiple texts
func (h *Handler) batchEmbeddingsWorkflow(r *http.Request, handler *Handler, input map[string]interface{}) (interface{}, error) {
	texts, ok := input["texts"].([]interface{})
	if !ok || len(texts) == 0 {
		return nil, fmt.Errorf("missing required parameter: texts (array)")
	}

	model, ok := input["model"].(string)
	if !ok {
		model = "nomic-ai/nomic-embed-text-v2-moe"
	}

	// Convert interface{} to []string
	textStrings := make([]string, 0)
	for _, t := range texts {
		if str, ok := t.(string); ok {
			textStrings = append(textStrings, str)
		}
	}

	if len(textStrings) == 0 {
		return nil, fmt.Errorf("no valid text strings in texts array")
	}

	embedReq := map[string]interface{}{
		"model": model,
		"input": textStrings,
	}

	embedBody, _ := json.Marshal(embedReq)
	embedHTTPReq, _ := http.NewRequest("POST", "/v1/embeddings", io.NopCloser(bytes.NewReader(embedBody)))
	embedHTTPReq.Header.Set("Content-Type", "application/json")

	embedResp := &responseCapture{}
	handler.ServeHTTP(embedResp, embedHTTPReq)

	var embedResult map[string]interface{}
	if err := json.Unmarshal(embedResp.body.Bytes(), &embedResult); err != nil {
		return nil, fmt.Errorf("failed to parse embedding response: %v", err)
	}

	return embedResult, nil
}

// Utility functions

// responseCapture captures HTTP response for reuse within workflows
type responseCapture struct {
	status int
	header http.Header
	body   bytes.Buffer
}

func (w *responseCapture) Header() http.Header {
	if w.header == nil {
		w.header = make(http.Header)
	}
	return w.header
}

func (w *responseCapture) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	return w.body.Write(b)
}

func (w *responseCapture) WriteHeader(statusCode int) {
	if w.status == 0 {
		w.status = statusCode
	}
}

// generateWorkflowID generates a unique workflow execution ID
func generateWorkflowID() string {
	return fmt.Sprintf("wf_%d", time.Now().UnixNano())
}