feat: implement OutputGuardrails for LLM response validation (#73)

Author: Raezil

agent.go

@@ -60,6 +60,7 @@ type Agent struct {
 	CodeChain *chain.UtcpChainClient
 
 	AllowUnsafeTools bool
+	Guardrails       *OutputGuardrails
 }
 
 // Options configure a new Agent.
@@ -77,6 +78,7 @@ type Options struct {
 	Shared            *memory.SharedSession
 	CodeChain         *chain.UtcpChainClient
 	AllowUnsafeTools  bool
+	Guardrails        *OutputGuardrails
 }
 
 // New creates an Agent with the provided options.
@@ -146,6 +148,7 @@ func New(opts Options) (*Agent, error) {
 		CodeMode:          opts.CodeMode,
 		CodeChain:         opts.CodeChain,
 		AllowUnsafeTools:  opts.AllowUnsafeTools,
+		Guardrails:        opts.Guardrails,
 	}
 
 	return a, nil
@@ -1214,7 +1217,17 @@ func (a *Agent) Generate(ctx context.Context, sessionID, userInput string) (any,
 	// 3. Chain Orchestrator (LLM decides a multi-step chain execution)
 	// -------------------------------------------------------------
 	if handled, output, err := a.codeChainOrchestrator(ctx, sessionID, userInput); handled {
-		return output, err
+		if err != nil {
+			return "", err
+		}
+		if a.Guardrails != nil {
+			validated, gErr := a.Guardrails.ValidateAndRepair(ctx, output)
+			if gErr != nil {
+				return "", gErr
+			}
+			return validated, nil
+		}
+		return output, nil
 	}
 	// ---------------------------------------------
 	// 4. TOOL ORCHESTRATOR (normal UTCP tools)
@@ -1270,7 +1283,17 @@ func (a *Agent) Generate(ctx context.Context, sessionID, userInput string) (any,
 		return "", err
 	}
 
-	a.storeMemory(sessionID, "assistant", fmt.Sprintf("%s", completion), nil)
+	finalText := fmt.Sprint(completion)
+	if a.Guardrails != nil {
+		validatedText, gErr := a.Guardrails.ValidateAndRepair(ctx, finalText)
+		if gErr != nil {
+			return "", gErr
+		}
+		finalText = validatedText
+		completion = finalText
+	}
+
+	a.storeMemory(sessionID, "assistant", finalText, nil)
 	return completion, nil
 }
 

agent_stream.go

@@ -73,6 +73,13 @@ func (a *Agent) GenerateStream(ctx context.Context, sessionID, userInput string)
 
 	// 3. Chain Orchestrator
 	if handled, output, err := a.codeChainOrchestrator(ctx, sessionID, userInput); handled {
+		if err == nil && a.Guardrails != nil {
+			validated, gErr := a.Guardrails.ValidateAndRepair(ctx, output)
+			if gErr != nil {
+				return immediateStream("", gErr)
+			}
+			output = validated
+		}
 		return immediateStream(output, err)
 	}
 
@@ -110,23 +117,51 @@ func (a *Agent) GenerateStream(ctx context.Context, sessionID, userInput string)
 
 	// Wrap the stream to intercept and store memory
 	outCh := make(chan models.StreamChunk)
-	go func() {
-		defer close(outCh)
-		var full strings.Builder
-		for chunk := range stream {
-			if chunk.Err != nil {
-				outCh <- chunk
+
+	if a.Guardrails != nil {
+		go func() {
+			defer close(outCh)
+			var full strings.Builder
+			for chunk := range stream {
+				if chunk.Err != nil {
+					outCh <- chunk
+					return
+				}
+				if chunk.Delta != "" {
+					full.WriteString(chunk.Delta)
+				}
+			}
+
+			finalText := full.String()
+			validatedText, gErr := a.Guardrails.ValidateAndRepair(ctx, finalText)
+			if gErr != nil {
+				outCh <- models.StreamChunk{Err: gErr, Done: true}
 				return
 			}
-			if chunk.Delta != "" {
-				full.WriteString(chunk.Delta)
+
+			// Stream out the validated text as one chunk
+			outCh <- models.StreamChunk{Delta: validatedText, FullText: validatedText, Done: true}
+			a.storeMemory(sessionID, "assistant", validatedText, nil)
+		}()
+	} else {
+		go func() {
+			defer close(outCh)
+			var full strings.Builder
+			for chunk := range stream {
+				if chunk.Err != nil {
+					outCh <- chunk
+					return
+				}
+				if chunk.Delta != "" {
+					full.WriteString(chunk.Delta)
+				}
+				outCh <- chunk
 			}
-			outCh <- chunk
-		}
-		// Store memory after completion
-		finalText := full.String()
-		a.storeMemory(sessionID, "assistant", finalText, nil)
-	}()
+			// Store memory after completion
+			finalText := full.String()
+			a.storeMemory(sessionID, "assistant", finalText, nil)
+		}()
+	}
 
 	return outCh, nil
 }

types.go

@@ -62,3 +62,41 @@ type AgentState struct {
 	JoinedSpaces []string                        `json:"joined_spaces,omitempty"`
 	Timestamp    time.Time                       `json:"timestamp"`
 }
+
+// SafetyPolicy defines an interface for validating LLM responses.
+type SafetyPolicy interface {
+	Validate(ctx context.Context, response string) error
+}
+
+// FormatEnforcer defines an interface for validating or repairing the format of LLM responses.
+type FormatEnforcer interface {
+	Enforce(ctx context.Context, response string) (string, error)
+}
+
+// OutputGuardrails holds the policy engines and formatting rules.
+type OutputGuardrails struct {
+	SafetyPolicies  []SafetyPolicy
+	FormatEnforcers []FormatEnforcer
+}
+
+// ValidateAndRepair applies safety checks and format enforcing to the response.
+func (g *OutputGuardrails) ValidateAndRepair(ctx context.Context, response string) (string, error) {
+	if g == nil {
+		return response, nil
+	}
+	for _, policy := range g.SafetyPolicies {
+		if err := policy.Validate(ctx, response); err != nil {
+			return "", err
+		}
+	}
+
+	final := response
+	var err error
+	for _, enforcer := range g.FormatEnforcers {
+		final, err = enforcer.Enforce(ctx, final)
+		if err != nil {
+			return "", err
+		}
+	}
+	return final, nil
+}