Production Agentic Search Optimization in Go

package cache

import (
	"context"
	"fmt"

	"github.com/redis/go-redis/v9"
	"github.com/cloudwego/eino-ext/components/retriever/redis"
)

// InitRedisRetriever establishes the connection and initializes the Eino Redis Retriever
func InitRedisRetriever(ctx context.Context) (*redis.Retriever, error) {
	// 1. Initialize go-redis client using Protocol 2 to be compatible with FT.SEARCH
	client := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "", // Enter password if any
		DB:       0,
		Protocol: 2, 
	})

	// Enable UnstableResp3 to support parsing complex response formats from RediSearch
	client.Options().UnstableResp3 = true

	// Check connection to the Redis Server
	if err := client.Ping(ctx).Err(); err != nil {
		return nil, fmt.Errorf("failed to connect to Redis: %w", err)
	}

	// 2. Initialize Retriever configuring eino-ext vector search
	retriever, err := redis.NewRetriever(ctx, &redis.RetrieverConfig{
		Client:       client,
		Index:        "semantic_cache_idx", // Vector index name on Redis
		VectorField:  "query_vector",       // Field storing the embedding
		EmbeddingKey: "query_text",         // Field storing the raw query
		TopK:         1,                    // Only fetch the most similar result
	})
	if err != nil {
		return nil, fmt.Errorf("failed to initialize Eino Redis Retriever: %w", err)
	}

	return retriever, nil
}

package routing

import (
	"context"
	"strings"

	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

// QueryState stores query information for routing decisions
type QueryState struct {
	Query     string
	IsComplex bool
}

// ModelRouterBranch executes dynamic routing based on the question state
var ModelRouterBranch = compose.NewGraphBranch(func(ctx context.Context, input *schema.Message) (string, error) {
	var nextNode string
	
	err := compose.ProcessState[*QueryState](ctx, func(ctx context.Context, state *QueryState) error {
		queryLower := strings.ToLower(state.Query)
		
		// Routing constraint: If the query is long (> 80 characters) or contains complex comparison/analysis keywords
		if len(state.Query) > 80 || 
			strings.Contains(queryLower, "compare") || 
			strings.Contains(queryLower, "analyze") || 
			strings.Contains(queryLower, "why") {
			
			state.IsComplex = true
			nextNode = "advanced_llm_node"
		} else {
			state.IsComplex = false
			nextNode = "cheap_llm_node"
		}
		return nil
	})
	
	return nextNode, err
}, map[string]bool{
	"cheap_llm_node":    true, // gpt-4o-mini node
	"advanced_llm_node": true, // Gemini 1.5 Pro node
})

package sse

import (
	"context"
	"fmt"
	"net/http"

	"github.com/cloudwego/eino/compose"
	"github.com/cloudwego/eino/schema"
)

// StreamSSEHandler processes the HTTP request and pushes data as Server-Sent Events
func StreamSSEHandler(w http.ResponseWriter, r *http.Request, runnable compose.Runnable[[]*schema.Message, *schema.StreamReader[*schema.Message]]) {
	// 1. Set required HTTP Headers for SSE
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Transfer-Encoding", "chunked")

	// Ensure the Web Server supports Streaming (Flusher)
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Browser or Server does not support Response Streaming", http.StatusInternalServerError)
		return
	}

	query := r.URL.Query().Get("q")
	if query == "" {
		http.Error(w, "Query parameter 'q' cannot be empty", http.StatusBadRequest)
		return
	}

	input := []*schema.Message{schema.UserMessage(query)}

	// 2. Activate Eino Stream to read data sequentially
	streamReader, err := runnable.Stream(r.Context(), input)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to initialize Stream: %v", err), http.StatusInternalServerError)
		return
	}
	// CRITICAL: Always Close streamReader to free the background Eino goroutines
	defer streamReader.Close() 

	// 3. Loop to receive data and push to Client
	for {
		msg, err := streamReader.Recv()
		if err != nil {
			// Receive EOF signal when the stream naturally concludes
			break
		}
		
		// Write data in standard SSE format (data: <content>\n\n)
		_, _ = fmt.Fprintf(w, "data: %s\n\n", msg.Content)
		flusher.Flush() // Push data out over the network immediately
	}

	// Send stream termination event for the Frontend to close the connection
	_, _ = fmt.Fprint(w, "event: done\ndata: [DONE]\n\n")
	flusher.Flush()
}

package telemetry

import (
	"context"

	"github.com/cloudwego/eino/callbacks"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

var tracer = otel.Tracer("eino-search-agent")

type spanCtxKey struct{}

// NewOTelCallbackHandler creates a custom event handler for monitoring
func NewOTelCallbackHandler() callbacks.Handler {
	return callbacks.NewHandlerBuilder().
		// Triggered when starting a Node in the graph
		OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
			// Start a new Span based on the respective Node's name
			ctx, span := tracer.Start(ctx, info.ComponentName, trace.WithSpanKind(trace.SpanKindInternal))
			span.SetAttributes(
				attribute.String("eino.component.type", string(info.ComponentType)),
				attribute.String("eino.component.name", info.ComponentName),
			)
			
			// Save Span to context so the next Node or the ending callback can access it
			return context.WithValue(ctx, spanCtxKey{}, span)
		}).
		// Triggered when the Node successfully completes processing
		OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
			if span, ok := ctx.Value(spanCtxKey{}).(trace.Span); ok {
				// Record token consumption according to OpenTelemetry Semantic Conventions
				if usage, ok := output.Config["token_usage"].(map[string]int); ok {
					span.SetAttributes(
						attribute.Int("gen_ai.usage.input_tokens", usage["input"]),
						attribute.Int("gen_ai.usage.output_tokens", usage["output"]),
					)
				}
				span.End() // Close Span
			}
			return ctx
		}).
		// Triggered when the Node encounters an error during execution
		OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
			if span, ok := ctx.Value(spanCtxKey{}).(trace.Span); ok {
				span.RecordError(err) // Mark error in Span
				span.End()
			}
			return ctx
		}).
		Build()
}