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 thiết lập kết nối và khởi tạo Eino Redis Retriever
func InitRedisRetriever(ctx context.Context) (*redis.Retriever, error) {
	// 1. Khởi tạo go-redis client dùng Protocol 2 để tương thích FT.SEARCH
	client := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "", // Điền mật khẩu nếu có
		DB:       0,
		Protocol: 2, 
	})

	// Kích hoạt UnstableResp3 để hỗ trợ phân tích định dạng phản hồi phức tạp từ RediSearch
	client.Options().UnstableResp3 = true

	// Kiểm tra kết nối tới Redis Server
	if err := client.Ping(ctx).Err(); err != nil {
		return nil, fmt.Errorf("không thể kết nối Redis: %w", err)
	}

	// 2. Khởi tạo Retriever cấu hình tìm kiếm vector của eino-ext
	retriever, err := redis.NewRetriever(ctx, &redis.RetrieverConfig{
		Client:       client,
		Index:        "semantic_cache_idx", // Tên chỉ mục vector trên Redis
		VectorField:  "query_vector",       // Trường lưu trữ embedding
		EmbeddingKey: "query_text",         // Trường lưu trữ câu hỏi thô
		TopK:         1,                    // Chỉ lấy kết quả tương tự nhất
	})
	if err != nil {
		return nil, fmt.Errorf("khởi tạo Eino Redis Retriever thất bại: %w", err)
	}

	return retriever, nil
}

package routing

import (
	"context"
	"strings"

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

// QueryState lưu trữ thông tin câu truy vấn để ra quyết định định tuyến
type QueryState struct {
	Query     string
	IsComplex bool
}

// ModelRouterBranch thực hiện định tuyến động dựa trên trạng thái câu hỏi
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)
		
		// Ràng buộc định tuyến: Nếu câu truy vấn dài (> 80 ký tự) hoặc chứa các từ khóa so sánh/phân tích phức tạp
		if len(state.Query) > 80 || 
			strings.Contains(queryLower, "so sánh") || 
			strings.Contains(queryLower, "phân tích") || 
			strings.Contains(queryLower, "tại sao") {
			
			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, // Node gpt-4o-mini
	"advanced_llm_node": true, // Node Gemini 1.5 Pro
})

package sse

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

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

// StreamSSEHandler xử lý yêu cầu HTTP và đẩy dữ liệu dạng Server-Sent Events
func StreamSSEHandler(w http.ResponseWriter, r *http.Request, runnable compose.Runnable[[]*schema.Message, *schema.StreamReader[*schema.Message]]) {
	// 1. Thiết lập các HTTP Header bắt buộc cho 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")

	// Đảm bảo Web Server hỗ trợ Streaming (Flusher)
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Trình duyệt hoặc máy chủ không hỗ trợ Response Streaming", http.StatusInternalServerError)
		return
	}

	query := r.URL.Query().Get("q")
	if query == "" {
		http.Error(w, "Tham số truy vấn 'q' không được trống", http.StatusBadRequest)
		return
	}

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

	// 2. Kích hoạt Eino Stream để đọc dữ liệu từng phần
	streamReader, err := runnable.Stream(r.Context(), input)
	if err != nil {
		http.Error(w, fmt.Sprintf("Không thể khởi tạo luồng Stream: %v", err), http.StatusInternalServerError)
		return
	}
	// QUAN TRỌNG: Luôn Close streamReader để giải phóng goroutine chạy ngầm của Eino
	defer streamReader.Close() 

	// 3. Vòng lặp nhận dữ liệu và đẩy xuống Client
	for {
		msg, err := streamReader.Recv()
		if err != nil {
			// Nhận tín hiệu EOF khi stream kết thúc tự nhiên
			break
		}
		
		// Ghi dữ liệu theo định dạng chuẩn SSE (data: <nội dung>\n\n)
		_, _ = fmt.Fprintf(w, "data: %s\n\n", msg.Content)
		flusher.Flush() // Đẩy dữ liệu ra mạng ngay lập tức
	}

	// Gửi sự kiện kết thúc luồng SSE để Frontend đóng kết nối
	_, _ = 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 tạo bộ xử lý sự kiện tùy biến phục vụ giám sát
func NewOTelCallbackHandler() callbacks.Handler {
	return callbacks.NewHandlerBuilder().
		// Kích hoạt khi bắt đầu chạy một Node trong đồ thị
		OnStartFn(func(ctx context.Context, info *callbacks.RunInfo, input callbacks.CallbackInput) context.Context {
			// Bắt đầu một Span mới dựa trên tên của Node tương ứng
			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),
			)
			
			// Lưu Span vào context để Node tiếp theo hoặc callback kết thúc có thể truy cập
			return context.WithValue(ctx, spanCtxKey{}, span)
		}).
		// Kích hoạt khi Node hoàn thành xử lý thành công
		OnEndFn(func(ctx context.Context, info *callbacks.RunInfo, output callbacks.CallbackOutput) context.Context {
			if span, ok := ctx.Value(spanCtxKey{}).(trace.Span); ok {
				// Ghi nhận số lượng token tiêu thụ theo đặc tả chuẩn 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() // Đóng Span
			}
			return ctx
		}).
		// Kích hoạt khi Node xảy ra lỗi trong quá trình thực thi
		OnErrorFn(func(ctx context.Context, info *callbacks.RunInfo, err error) context.Context {
			if span, ok := ctx.Value(spanCtxKey{}).(trace.Span); ok {
				span.RecordError(err) // Đánh dấu lỗi trong Span
				span.End()
			}
			return ctx
		}).
		Build()
}