internal/workflow/rag_workflow.go

// Package workflow provides Temporal.io workflow definitions for RAG pipelines
package workflow

import (
	"context"
	"time"

	"go.temporal.io/sdk/activity"
	"go.temporal.io/sdk/temporal"
	"go.temporal.io/sdk/workflow"
)

// RAGRequest represents a RAG pipeline request
type RAGRequest struct {
	Query       string
	SessionID   string
	UserID      string
	TopK        int
	UseVector   bool
	UseKeyword  bool
	UseGraph    bool
	Temperature float32
	MaxTokens   int
}

// RAGResponse represents a RAG pipeline response
type RAGResponse struct {
	Answer     string
	Sources    []Source
	Confidence float32
	Metadata   RAGMetadata
}

// Source represents a retrieved source
type Source struct {
	ID       string
	Title    string
	Content  string
	Score    float32
	Type     string // "vector", "keyword", "graph"
	Metadata map[string]string
}

// RAGMetadata contains pipeline execution metadata
type RAGMetadata struct {
	TotalLatencyMs     int64
	EmbeddingLatencyMs int64
	RetrievalLatencyMs int64
	GenerationLatencyMs int64
	TokensUsed         int
	ChunksRetrieved    int
}

// ActivityDependencies contains all activity dependencies
type ActivityDependencies struct {
	EmbeddingClient interface {
		Generate(ctx context.Context, text string) ([]float32, error)
	}
	VectorStore interface {
		Search(ctx context.Context, embedding []float32, topK int, filters map[string]interface{}) ([]Source, error)
	}
	KeywordEngine interface {
		Search(ctx context.Context, query string, topK int) ([]Source, error)
	}
	GraphStore interface {
		Search(ctx context.Context, query string, embedding []float32, topK int) ([]Source, error)
	}
	Guardrails interface {
		ValidateInput(ctx context.Context, input string) (bool, string, error)
		ValidateOutput(ctx context.Context, input, output string) (bool, string, error)
	}
	LLMClient interface {
		Generate(ctx context.Context, prompt string, options GenerateOptions) (string, int, error)
	}
}

// GenerateOptions for LLM generation
type GenerateOptions struct {
	Temperature float32
	MaxTokens   int
	Context     []Source
}

// RAGWorkflow orchestrates the complete RAG pipeline with durable execution
func RAGWorkflow(ctx workflow.Context, req RAGRequest) (*RAGResponse, error) {
	logger := workflow.GetLogger(ctx)
	logger.Info("starting RAG workflow", "query", req.Query, "session_id", req.SessionID)

	// Configure activity options with retries
	ao := workflow.ActivityOptions{
		StartToCloseTimeout: 60 * time.Second,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    time.Second,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    3,
		},
	}
	ctx = workflow.WithActivityOptions(ctx, ao)

	startTime := workflow.Now(ctx)
	var metadata RAGMetadata

	// Step 1: Input validation (guardrails)
	var inputValid bool
	var inputReason string
	err := workflow.ExecuteActivity(ctx, ValidateInputActivity, req.Query).Get(ctx, &inputValid)
	if err != nil {
		logger.Warn("input validation failed, continuing", "error", err)
		inputValid = true // Fail open if guardrails unavailable
	}
	if !inputValid {
		return &RAGResponse{
			Answer: "I'm sorry, but I cannot process this request. " + inputReason,
			Metadata: RAGMetadata{
				TotalLatencyMs: workflow.Now(ctx).Sub(startTime).Milliseconds(),
			},
		}, nil
	}

	// Step 2: Generate embeddings
	embeddingStart := workflow.Now(ctx)
	var embedding []float32
	err = workflow.ExecuteActivity(ctx, GenerateEmbeddingActivity, req.Query).Get(ctx, &embedding)
	if err != nil {
		return nil, err
	}
	metadata.EmbeddingLatencyMs = workflow.Now(ctx).Sub(embeddingStart).Milliseconds()

	// Step 3: Parallel retrieval (vector, keyword, graph)
	retrievalStart := workflow.Now(ctx)
	var allSources []Source

	// Use futures for parallel execution
	var vectorFuture, keywordFuture, graphFuture workflow.Future

	if req.UseVector {
		vectorFuture = workflow.ExecuteActivity(ctx, VectorSearchActivity, VectorSearchInput{
			Embedding: embedding,
			TopK:      req.TopK,
		})
	}

	if req.UseKeyword {
		keywordFuture = workflow.ExecuteActivity(ctx, KeywordSearchActivity, KeywordSearchInput{
			Query: req.Query,
			TopK:  req.TopK,
		})
	}

	if req.UseGraph {
		graphFuture = workflow.ExecuteActivity(ctx, GraphSearchActivity, GraphSearchInput{
			Query:     req.Query,
			Embedding: embedding,
			TopK:      req.TopK,
		})
	}

	// Collect results
	if vectorFuture != nil {
		var vectorSources []Source
		if err := vectorFuture.Get(ctx, &vectorSources); err != nil {
			logger.Warn("vector search failed", "error", err)
		} else {
			allSources = append(allSources, vectorSources...)
		}
	}

	if keywordFuture != nil {
		var keywordSources []Source
		if err := keywordFuture.Get(ctx, &keywordSources); err != nil {
			logger.Warn("keyword search failed", "error", err)
		} else {
			allSources = append(allSources, keywordSources...)
		}
	}

	if graphFuture != nil {
		var graphSources []Source
		if err := graphFuture.Get(ctx, &graphSources); err != nil {
			logger.Warn("graph search failed", "error", err)
		} else {
			allSources = append(allSources, graphSources...)
		}
	}

	metadata.RetrievalLatencyMs = workflow.Now(ctx).Sub(retrievalStart).Milliseconds()
	metadata.ChunksRetrieved = len(allSources)

	// Step 4: Rank and deduplicate sources
	var rankedSources []Source
	err = workflow.ExecuteActivity(ctx, RankSourcesActivity, RankSourcesInput{
		Sources: allSources,
		Query:   req.Query,
		TopK:    req.TopK,
	}).Get(ctx, &rankedSources)
	if err != nil {
		rankedSources = allSources // Use unranked if ranking fails
	}

	// Step 5: Generate response
	generationStart := workflow.Now(ctx)
	var generateResult GenerateResult
	err = workflow.ExecuteActivity(ctx, GenerateResponseActivity, GenerateInput{
		Query:       req.Query,
		Sources:     rankedSources,
		Temperature: req.Temperature,
		MaxTokens:   req.MaxTokens,
	}).Get(ctx, &generateResult)
	if err != nil {
		return nil, err
	}
	metadata.GenerationLatencyMs = workflow.Now(ctx).Sub(generationStart).Milliseconds()
	metadata.TokensUsed = generateResult.TokensUsed

	// Step 6: Output validation (guardrails)
	var outputValid bool
	err = workflow.ExecuteActivity(ctx, ValidateOutputActivity, ValidateOutputInput{
		Input:  req.Query,
		Output: generateResult.Content,
	}).Get(ctx, &outputValid)
	if err != nil {
		logger.Warn("output validation failed, continuing", "error", err)
		outputValid = true
	}
	if !outputValid {
		generateResult.Content = "I apologize, but I cannot provide this response due to safety guidelines."
	}

	metadata.TotalLatencyMs = workflow.Now(ctx).Sub(startTime).Milliseconds()

	return &RAGResponse{
		Answer:     generateResult.Content,
		Sources:    rankedSources,
		Confidence: calculateConfidence(rankedSources),
		Metadata:   metadata,
	}, nil
}

func calculateConfidence(sources []Source) float32 {
	if len(sources) == 0 {
		return 0.0
	}
	var total float32
	for _, s := range sources {
		total += s.Score
	}
	return total / float32(len(sources))
}

// Activity input/output types

type VectorSearchInput struct {
	Embedding []float32
	TopK      int
}

type KeywordSearchInput struct {
	Query string
	TopK  int
}

type GraphSearchInput struct {
	Query     string
	Embedding []float32
	TopK      int
}

type RankSourcesInput struct {
	Sources []Source
	Query   string
	TopK    int
}

type GenerateInput struct {
	Query       string
	Sources     []Source
	Temperature float32
	MaxTokens   int
}

type GenerateResult struct {
	Content    string
	TokensUsed int
}

type ValidateOutputInput struct {
	Input  string
	Output string
}

// Activity implementations (stubs - implemented in activities.go)

func ValidateInputActivity(ctx context.Context, input string) (bool, error) {
	return true, nil // Implemented in activities.go
}

func GenerateEmbeddingActivity(ctx context.Context, text string) ([]float32, error) {
	logger := activity.GetLogger(ctx)
	logger.Info("generating embedding", "text_length", len(text))
	return nil, nil // Implemented with actual client
}

func VectorSearchActivity(ctx context.Context, input VectorSearchInput) ([]Source, error) {
	return nil, nil
}

func KeywordSearchActivity(ctx context.Context, input KeywordSearchInput) ([]Source, error) {
	return nil, nil
}

func GraphSearchActivity(ctx context.Context, input GraphSearchInput) ([]Source, error) {
	return nil, nil
}

func RankSourcesActivity(ctx context.Context, input RankSourcesInput) ([]Source, error) {
	return input.Sources, nil // Simple passthrough
}

func GenerateResponseActivity(ctx context.Context, input GenerateInput) (*GenerateResult, error) {
	return nil, nil
}

func ValidateOutputActivity(ctx context.Context, input ValidateOutputInput) (bool, error) {
	return true, nil
}