Week 2 Day 9：Embedding + Vector Store - 苏格拉底教学#

// internal/rag/embedder.go
package rag

import (
	"context"
	"fmt"

	openai "github.com/sashabaranov/go-openai"
)

type Embedder struct {
	client *openai.Client
	model  string
}

func NewEmbedder(apiKey string) *Embedder {
	return &Embedder{
		client: openai.NewClient(apiKey),
		model:  string(openai.SmallEmbedding3), // text-embedding-3-small
	}
}

// EmbedOne 对单个文本生成 embedding
func (e *Embedder) EmbedOne(ctx context.Context, text string) ([]float32, error) {
	resp, err := e.client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
		Model: openai.EmbeddingModel(e.model),
		Input: []string{text},
	})
	if err != nil {
		return nil, fmt.Errorf("embedding failed: %w", err)
	}

	if len(resp.Data) == 0 {
		return nil, fmt.Errorf("no embedding returned")
	}

	return resp.Data[0].Embedding, nil
}

// internal/rag/embedder.go（续）

import (
	"context"
	"fmt"
	"log/slog"
	"sync"

	openai "github.com/sashabaranov/go-openai"
	"golang.org/x/sync/errgroup"
)

const (
	batchSize      = 20 // OpenAI 每次最多 100 个，保守用 20
	maxConcurrency = 5  // 最多同时 5 个并发请求
)

// EmbedBatch 并发批量 embedding
// errgroup 负责错误收集，semaphore 控制并发度
func (e *Embedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error) {
	results := make([][]float32, len(texts))

	// semaphore：令牌桶，限制最大并发数
	sem := make(chan struct{}, maxConcurrency)

	g, gctx := errgroup.WithContext(ctx)
	var mu sync.Mutex

	// 分批处理
	for batchStart := 0; batchStart < len(texts); batchStart += batchSize {
		start := batchStart
		end := start + batchSize
		if end > len(texts) {
			end = len(texts)
		}
		batch := texts[start:end]

		g.Go(func() error {
			// 获取令牌（若已有 5 个并发则阻塞等待）
			sem <- struct{}{}
			defer func() { <-sem }()

			slog.Info("embedding batch", "start", start, "end", end, "size", len(batch))

			resp, err := e.client.CreateEmbeddings(gctx, openai.EmbeddingRequest{
				Model: openai.EmbeddingModel(e.model),
				Input: batch,
			})
			if err != nil {
				return fmt.Errorf("batch [%d:%d] failed: %w", start, end, err)
			}

			// 写回结果，加锁防止竞态
			mu.Lock()
			for i, data := range resp.Data {
				results[start+i] = data.Embedding
			}
			mu.Unlock()

			return nil
		})
	}

	// 等待所有 goroutine 完成，有任何错误则返回
	if err := g.Wait(); err != nil {
		return nil, err
	}

	return results, nil
}

// internal/rag/vectorstore.go
package rag

import (
	"context"
	"fmt"

	"github.com/google/uuid"
	"github.com/qdrant/go-client/qdrant"
)

const (
	collectionName = "knowledge_base"
	vectorSize     = 1536 // text-embedding-3-small 维度
)

type VectorStore struct {
	client         *qdrant.Client
	collectionName string
}

func NewVectorStore(host string, port int) (*VectorStore, error) {
	client, err := qdrant.NewClient(&qdrant.Config{
		Host: host,
		Port: port,
	})
	if err != nil {
		return nil, fmt.Errorf("connect qdrant: %w", err)
	}

	return &VectorStore{
		client:         client,
		collectionName: collectionName,
	}, nil
}

// EnsureCollection 确保 collection 存在（幂等）
func (vs *VectorStore) EnsureCollection(ctx context.Context) error {
	exists, err := vs.client.CollectionExists(ctx, vs.collectionName)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}

	return vs.client.CreateCollection(ctx, &qdrant.CreateCollection{
		CollectionName: vs.collectionName,
		VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{
			Size:     vectorSize,
			Distance: qdrant.Distance_Cosine,
		}),
	})
}

// chunkIDToUUID 把 chunk ID 稳定转换为 UUID（不丢失原始信息）
func chunkIDToUUID(chunkID string) string {
	uid := uuid.NewSHA1(uuid.NameSpaceURL, []byte(chunkID))
	return uid.String()
}

// UpsertChunks 批量写入 chunks 和对应向量
func (vs *VectorStore) UpsertChunks(ctx context.Context, chunks []Chunk, vectors [][]float32) error {
	if len(chunks) != len(vectors) {
		return fmt.Errorf("chunks(%d) and vectors(%d) length mismatch", len(chunks), len(vectors))
	}

	points := make([]*qdrant.PointStruct, len(chunks))
	for i, chunk := range chunks {
		uid := chunkIDToUUID(chunk.ID)
		points[i] = &qdrant.PointStruct{
			Id:      qdrant.NewIDWithUUID(uid),
			Vectors: qdrant.NewVectors(vectors[i]...),
			Payload: map[string]*qdrant.Value{
				"chunk_id":       qdrant.NewValueString(chunk.ID),
				"text":           qdrant.NewValueString(chunk.Text),
				"doc_id":         qdrant.NewValueString(chunk.DocID),
				"title":          qdrant.NewValueString(chunk.Title),
				"source":         qdrant.NewValueString(chunk.Source),
				"department":     qdrant.NewValueString(chunk.Metadata["department"]),
				"classification": qdrant.NewValueString(chunk.Metadata["classification"]),
				"chunk_index":    qdrant.NewValueString(chunk.Metadata["chunk_index"]),
			},
		}
	}

	_, err := vs.client.Upsert(ctx, &qdrant.UpsertPoints{
		CollectionName: vs.collectionName,
		Points:         points,
	})
	return err
}

// Search 向量相似度搜索
func (vs *VectorStore) Search(ctx context.Context, queryVec []float32, topK uint64) ([]*qdrant.ScoredPoint, error) {
	results, err := vs.client.Query(ctx, &qdrant.QueryPoints{
		CollectionName: vs.collectionName,
		Query:          qdrant.NewQuery(queryVec...),
		Limit:          &topK,
		WithPayload:    qdrant.NewWithPayload(true),
	})
	if err != nil {
		return nil, fmt.Errorf("vector search: %w", err)
	}

	return results, nil
}

// cmd/indexer/main.go
package main

import (
	"context"
	"encoding/json"
	"log/slog"
	"os"
	"time"

	"github.com/yourname/agent-runtime/internal/rag"
)

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()

	apiKey := os.Getenv("OPENAI_API_KEY")
	if apiKey == "" {
		slog.Error("OPENAI_API_KEY not set")
		os.Exit(1)
	}

	// 1. 读取 Day 8 生成的 chunks.jsonl
	chunks, err := loadChunks("evals/chunks.jsonl")
	if err != nil {
		slog.Error("load chunks failed", "error", err)
		os.Exit(1)
	}
	slog.Info("loaded chunks", "count", len(chunks))

	// 2. 提取文本，用 title + text 一起 embed，增强检索精度
	texts := make([]string, len(chunks))
	for i, c := range chunks {
		texts[i] = c.Title + "\n\n" + c.Text
	}

	// 3. 并发批量 embedding
	embedder := rag.NewEmbedder(apiKey)
	slog.Info("starting batch embedding", "chunks", len(texts))

	vectors, err := embedder.EmbedBatch(ctx, texts)
	if err != nil {
		slog.Error("embedding failed", "error", err)
		os.Exit(1)
	}
	slog.Info("embedding complete", "vectors", len(vectors))

	// 4. 初始化 Qdrant，确保 collection 存在
	vs, err := rag.NewVectorStore("localhost", 6334)
	if err != nil {
		slog.Error("vectorstore init failed", "error", err)
		os.Exit(1)
	}

	if err := vs.EnsureCollection(ctx); err != nil {
		slog.Error("ensure collection failed", "error", err)
		os.Exit(1)
	}

	// 5. 写入向量
	if err := vs.UpsertChunks(ctx, chunks, vectors); err != nil {
		slog.Error("upsert failed", "error", err)
		os.Exit(1)
	}

	slog.Info("indexing complete", "total_chunks", len(chunks))
}

func loadChunks(path string) ([]rag.Chunk, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var chunks []rag.Chunk
	decoder := json.NewDecoder(file)
	for decoder.More() {
		var c rag.Chunk
		if err := decoder.Decode(&c); err != nil {
			return nil, err
		}
		chunks = append(chunks, c)
	}
	return chunks, nil
}

// 正确用法：g.Go 内必须捕获所有错误
g, gctx := errgroup.WithContext(parentCtx)

g.Go(func() error {
    // 注意：使用 gctx 而不是 parentCtx
    // 若其他 goroutine 出错，gctx 会被取消
    result, err := doWork(gctx)
    if err != nil {
        return err // errgroup 会收集第一个 error
    }
    // 处理 result...
    return nil
})

// g.Wait() 会：
// 1. 等所有 goroutine 完成
// 2. 返回第一个非 nil error
// 3. 同时取消 gctx（通知其他 goroutine 停止）
if err := g.Wait(); err != nil {
    return err
}

g.Go(func() error {
    for {
        select {
        case <-gctx.Done():
            return gctx.Err() // 收到取消信号，主动退出
        default:
            // 继续工作
        }
    }
})

// 1. query 端和 document 端必须用同一模型
// 错误：document 用 ada-002，query 用 3-small → 向量空间不兼容
// 正确：统一用 text-embedding-3-small

// 2. 加前缀可提升部分模型的精度（E5 系列推荐）
func prepareText(text string, isQuery bool) string {
    if isQuery {
        return "query: " + text
    }
    return "passage: " + text
}
// OpenAI 3-small 不强制要求前缀，但可以尝试

// 3. 批量请求远优于单个请求
// 串行 1000 次：约 200 秒
// 批量（20/批）+ 并发 5：约 2 秒（100x 加速）

// 4. 考虑 embedding 缓存（相同文本不重复调用）
type CachedEmbedder struct {
    inner *Embedder
    cache map[string][]float32
    mu    sync.RWMutex
}

// cmd/search/main.go
// 测试：把 query embed 后搜索，人工判断返回 chunks 是否相关
func main() {
    query := "如何重置密码"
    // 1. embed query
    vec, _ := embedder.EmbedOne(ctx, query)
    // 2. 搜索 top-5
    results, _ := vs.Search(ctx, vec, 5)
    // 3. 打印返回 chunk 的 text，人工判断相关性
    for _, r := range results {
        fmt.Printf("score=%.4f text=%s\n",
            r.Score, r.Payload["text"].GetStringValue())
    }
}

import "github.com/cenkalti/backoff/v4"

operation := func() error {
    resp, err := e.client.CreateEmbeddings(ctx, req)
    if err != nil {
        return err
    }
    // 处理 resp...
    return nil
}

err := backoff.Retry(operation, backoff.WithMaxRetries(
    backoff.NewExponentialBackOff(), 3,
))

return vs.client.CreateCollection(ctx, &qdrant.CreateCollection{
    CollectionName: vs.collectionName,
    VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{
        Size:     vectorSize,
        Distance: qdrant.Distance_Cosine,
        HnswConfig: &qdrant.HnswConfigDiff{
            M:           qdrant.PtrOf(uint64(16)),  // 每个节点的连接数，越大精度越高，内存越多
            EfConstruct: qdrant.PtrOf(uint64(100)), // 构建时的搜索宽度
        },
    }),
})

// 文档更新时：重新生成 chunk 和 embedding，用相同 UUID upsert 即可
// 注意：chunk ID → UUID 的映射必须保持不变（用同一个 hash 函数）

// LeetCode 102: Binary Tree Level Order Traversal
// 返回每层节点值的列表
func levelOrder(root *TreeNode) [][]int {
    if root == nil {
        return nil
    }

    var result [][]int
    queue := []*TreeNode{root}

    for len(queue) > 0 {
        levelSize := len(queue) // 记录当前层的节点数（关键！）
        var level []int

        for i := 0; i < levelSize; i++ {
            node := queue[0]
            queue = queue[1:]

            level = append(level, node.Val)

            if node.Left != nil {
                queue = append(queue, node.Left)
            }
            if node.Right != nil {
                queue = append(queue, node.Right)
            }
        }

        result = append(result, level)
    }

    return result
}

// 只需在追加 level 时按层奇偶决定是否 reverse
if len(result)%2 == 1 {
    // 反转 level
    for i, j := 0, len(level)-1; i < j; i, j = i+1, j-1 {
        level[i], level[j] = level[j], level[i]
    }
}
result = append(result, level)

// LeetCode 98: Validate Binary Search Tree
func isValidBST(root *TreeNode) bool {
    return validate(root, nil, nil)
}

// 用 min/max 边界传递约束，而不是只比较父子节点
// 错误做法：只看 node.Val > node.Left.Val（遗漏跨层约束）
// 正确做法：整个左子树的所有节点都必须 < 当前节点
func validate(node *TreeNode, min, max *int) bool {
    if node == nil {
        return true
    }

    if min != nil && node.Val <= *min {
        return false
    }
    if max != nil && node.Val >= *max {
        return false
    }

    // 左子树：上界更新为当前节点值
    // 右子树：下界更新为当前节点值
    return validate(node.Left, min, &node.Val) &&
        validate(node.Right, &node.Val, max)
}

// LeetCode 230: Kth Smallest Element in a BST
// 思路：中序遍历 BST 得到升序序列，第 k 个即答案
func kthSmallest(root *TreeNode, k int) int {
    count := 0
    result := 0

    var inorder func(node *TreeNode)
    inorder = func(node *TreeNode) {
        if node == nil || count >= k {
            return
        }

        inorder(node.Left)

        count++
        if count == k {
            result = node.Val
            return
        }

        inorder(node.Right)
    }

    inorder(root)
    return result
}

// 进阶：如果频繁查询 kth smallest，怎么优化到 O(log n)？
// 思路：在每个节点额外记录左子树大小
type AugmentedNode struct {
    Val       int
    Left      *AugmentedNode
    Right     *AugmentedNode
    LeftCount int // 左子树节点总数
}

func kthSmallestAugmented(root *AugmentedNode, k int) int {
    if root == nil {
        return -1
    }
    leftCount := root.LeftCount
    if k == leftCount+1 {
        return root.Val // 当前节点就是第 k 小
    } else if k <= leftCount {
        return kthSmallestAugmented(root.Left, k)
    } else {
        return kthSmallestAugmented(root.Right, k-leftCount-1)
    }
}