Go 与 AI：构建智能应用

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// Client 是 AI API 的通用客户端
type Client struct {
	apiKey     string
	baseURL    string
	httpClient *http.Client
	model      string
}

// ClientOption 是客户端配置选项
type ClientOption func(*Client)

func WithHTTPClient(c *http.Client) ClientOption {
	return func(client *Client) {
		client.httpClient = c
	}
}

func WithTimeout(d time.Duration) ClientOption {
	return func(client *Client) {
		client.httpClient.Timeout = d
	}
}

func WithBaseURL(url string) ClientOption {
	return func(client *Client) {
		client.baseURL = url
	}
}

// NewOpenAIClient 创建 OpenAI 兼容的客户端
func NewOpenAIClient(apiKey string, opts ...ClientOption) *Client {
	c := &Client{
		apiKey:  apiKey,
		baseURL: "https://api.openai.com/v1",
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
		},
		model: "gpt-4o",
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// NewClaudeClient 创建 Anthropic Claude 客户端
func NewClaudeClient(apiKey string, opts ...ClientOption) *Client {
	c := &Client{
		apiKey:  apiKey,
		baseURL: "https://api.anthropic.com/v1",
		httpClient: &http.Client{
			Timeout: 120 * time.Second,
		},
		model: "claude-sonnet-4-20250514",
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// Message 表示一条对话消息
type Message struct {
	Role       string     `json:"role"`
	Content    string     `json:"content,omitempty"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
}

// ToolCall 表示模型发起的工具调用
type ToolCall struct {
	ID       string       `json:"id"`
	Type     string       `json:"type"`
	Function FunctionCall `json:"function"`
}

// FunctionCall 表示函数调用的详情
type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// ChatRequest 是聊天请求
type ChatRequest struct {
	Model       string     `json:"model"`
	Messages    []Message  `json:"messages"`
	Temperature float64    `json:"temperature,omitempty"`
	MaxTokens   int        `json:"max_tokens,omitempty"`
	Stream      bool       `json:"stream,omitempty"`
	Tools       []Tool     `json:"tools,omitempty"`
	Stop        []string   `json:"stop,omitempty"`
}

// Tool 定义一个可供模型调用的工具
type Tool struct {
	Type     string         `json:"type"`
	Function FunctionDef    `json:"function"`
}

// FunctionDef 定义函数的参数和描述
type FunctionDef struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters"`
}

// ChatResponse 是聊天响应
type ChatResponse struct {
	ID      string   `json:"id"`
	Choices []Choice `json:"choices"`
	Usage   Usage    `json:"usage"`
}

// Choice 是响应中的一个选择
type Choice struct {
	Index        int     `json:"index"`
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

// Usage 记录 token 使用量
type Usage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// Chat 发送聊天请求并获取响应
func (c *Client) Chat(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
	if req.Model == "" {
		req.Model = c.model
	}

	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("序列化请求失败: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("发送请求失败: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("API 错误 (status=%d): %s", resp.StatusCode, string(respBody))
	}

	var chatResp ChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
		return nil, fmt.Errorf("解析响应失败: %w", err)
	}

	return &chatResp, nil
}

package main

import (
	"context"
	"fmt"
	"log"
	"os"

	"example.com/aiapp/ai"
)

func main() {
	client := ai.NewOpenAIClient(os.Getenv("OPENAI_API_KEY"))

	resp, err := client.Chat(context.Background(), &ai.ChatRequest{
		Messages: []ai.Message{
			{Role: "system", Content: "你是一个专业的 Go 语言导师，回答简洁且有深度。"},
			{Role: "user", Content: "Go 语言的 goroutine 和 Java 的线程有什么本质区别？"},
		},
		Temperature: 0.7,
		MaxTokens:   1000,
	})
	if err != nil {
		log.Fatalf("调用失败: %v", err)
	}

	fmt.Println(resp.Choices[0].Message.Content)
	fmt.Printf("\n--- Token 使用量 ---\n")
	fmt.Printf("Prompt: %d, Completion: %d, Total: %d\n",
		resp.Usage.PromptTokens, resp.Usage.CompletionTokens, resp.Usage.TotalTokens)
}

package ai

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// StreamChunk 表示流式响应中的一个片段
type StreamChunk struct {
	ID      string        `json:"id"`
	Choices []StreamChoice `json:"choices"`
	Usage   *Usage        `json:"usage,omitempty"`
}

// StreamChoice 是流式片段中的选择
type StreamChoice struct {
	Index        int          `json:"index"`
	Delta        DeltaMessage `json:"delta"`
	FinishReason *string      `json:"finish_reason"`
}

// DeltaMessage 是增量消息内容
type DeltaMessage struct {
	Role      string     `json:"role,omitempty"`
	Content   string     `json:"content,omitempty"`
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

// StreamHandler 处理每个流式片段的回调函数
type StreamHandler func(chunk StreamChunk) error

// ChatStream 以流式方式获取 AI 响应
func (c *Client) ChatStream(ctx context.Context, req *ChatRequest, handler StreamHandler) error {
	req.Stream = true
	if req.Model == "" {
		req.Model = c.model
	}

	body, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("序列化请求失败: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("创建请求失败: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
	httpReq.Header.Set("Accept", "text/event-stream")

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return fmt.Errorf("发送请求失败: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("API 错误 (status=%d): %s", resp.StatusCode, string(respBody))
	}

	// 解析 SSE 流
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()

		// SSE 格式：以 "data: " 开头
		if !strings.HasPrefix(line, "data: ") {
			continue
		}

		data := strings.TrimPrefix(line, "data: ")

		// [DONE] 表示流结束
		if data == "[DONE]" {
			break
		}

		var chunk StreamChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			// 跳过解析失败的行
			continue
		}

		if err := handler(chunk); err != nil {
			return fmt.Errorf("处理片段失败: %w", err)
		}
	}

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("读取流失败: %w", err)
	}

	return nil
}

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"example.com/aiapp/ai"
)

func main() {
	client := ai.NewOpenAIClient(os.Getenv("OPENAI_API_KEY"))

	messages := []ai.Message{
		{Role: "system", Content: "你是一个创意故事作家。请用生动的语言写一个简短的故事。"},
		{Role: "user", Content: "写一个关于一个程序员穿越到古代的故事，要幽默。"},
	}

	fmt.Println("🤖 AI 正在创作中...")
	fmt.Println(strings.Repeat("─", 60))

	start := time.Now()
	tokenCount := 0

	// 收集完整响应用于后续对话
	var fullResponse strings.Builder

	err := client.ChatStream(context.Background(), &ai.ChatRequest{
		Messages:    messages,
		Temperature: 0.9,
		MaxTokens:   2000,
	}, func(chunk ai.StreamChunk) error {
		if len(chunk.Choices) > 0 {
			content := chunk.Choices[0].Delta.Content
			if content != "" {
				fmt.Print(content)
				fullResponse.WriteString(content)
				tokenCount++
			}

			// 检查是否结束
			if chunk.Choices[0].FinishReason != nil {
				fmt.Println()
			}
		}

		// 打印最终的 token 使用统计
		if chunk.Usage != nil {
			fmt.Println(strings.Repeat("─", 60))
			fmt.Printf("📊 Token 统计: Prompt=%d, Completion=%d, Total=%d\n",
				chunk.Usage.PromptTokens, chunk.Usage.CompletionTokens, chunk.Usage.TotalTokens)
		}

		return nil
	})

	if err != nil {
		log.Fatalf("流式调用失败: %v", err)
	}

	elapsed := time.Since(start)
	fmt.Printf("⏱  耗时: %v\n", elapsed)
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"strconv"
	"time"

	"example.com/aiapp/ai"
)

// 定义工具函数
var tools = []ai.Tool{
	{
		Type: "function",
		Function: ai.FunctionDef{
			Name:        "get_weather",
			Description: "获取指定城市的当前天气信息，包括温度、湿度和天气状况",
			Parameters: json.RawMessage(`{
				"type": "object",
				"properties": {
					"city": {
						"type": "string",
						"description": "城市名称，例如 '北京' 或 '上海'"
					},
					"unit": {
						"type": "string",
						"enum": ["celsius", "fahrenheit"],
						"description": "温度单位"
					}
				},
				"required": ["city"]
			}`),
		},
	},
	{
		Type: "function",
		Function: ai.FunctionDef{
			Name:        "calculate",
			Description: "执行数学计算，支持基本运算和常用数学函数",
			Parameters: json.RawMessage(`{
				"type": "object",
				"properties": {
					"expression": {
						"type": "string",
						"description": "数学表达式，例如 'sqrt(144) + 2^10'"
					}
				},
				"required": ["expression"]
			}`),
		},
	},
	{
		Type: "function",
		Function: ai.FunctionDef{
			Name:        "search_database",
			Description: "在数据库中搜索用户信息",
			Parameters: json.RawMessage(`{
				"type": "object",
				"properties": {
					"query": {
						"type": "string",
						"description": "搜索关键词"
					},
					"limit": {
						"type": "integer",
						"description": "返回结果数量上限",
						"default": 5
					}
				},
				"required": ["query"]
			}`),
		},
	},
}

// 模拟天气 API
func getWeather(city, unit string) string {
	// 实际应用中这里会调用真实的天气 API
	weatherData := map[string]map[string]interface{}{
		"北京": {"temp": 28, "humidity": 45, "condition": "晴", "wind": "东北风 3级"},
		"上海": {"temp": 32, "humidity": 78, "condition": "多云", "wind": "东南风 2级"},
		"广州": {"temp": 35, "humidity": 85, "condition": "雷阵雨", "wind": "南风 1级"},
	}

	data, ok := weatherData[city]
	if !ok {
		return fmt.Sprintf("抱歉，未找到 %s 的天气数据", city)
	}

	tempUnit := "°C"
	if unit == "fahrenheit" {
		data["temp"] = data["temp"].(int)*9/5 + 32
		tempUnit = "°F"
	}

	return fmt.Sprintf("📍 %s 当前天气：%s，温度 %d%s，湿度 %d%%，%s",
		city, data["condition"], data["temp"], tempUnit, data["humidity"], data["wind"])
}

// 模拟数学计算
func calculate(expression string) string {
	// 简单的表达式求值（实际应用中应使用安全的表达式解析器）
	result, err := strconv.ParseFloat(expression, 64)
	if err != nil {
		// 尝试一些预定义的计算
		switch expression {
		case "sqrt(144)":
			return fmt.Sprintf("sqrt(144) = %.2f", math.Sqrt(144))
		case "2^10":
			return fmt.Sprintf("2^10 = %d", int(math.Pow(2, 10)))
		}
		return fmt.Sprintf("无法计算表达式: %s", expression)
	}
	return fmt.Sprintf("%s = %.4f", expression, result)
}

// 模拟数据库搜索
func searchDatabase(query string, limit int) string {
	// 模拟数据库结果
	users := []map[string]string{
		{"name": "张三", "email": "zhangsan@example.com", "role": "工程师"},
		{"name": "李四", "email": "lisi@example.com", "role": "产品经理"},
		{"name": "王五", "email": "wangwu@example.com", "role": "设计师"},
	}

	results := users
	if limit < len(results) {
		results = results[:limit]
	}

	jsonData, _ := json.MarshalIndent(results, "", "  ")
	return string(jsonData)
}

// 执行工具调用并返回结果
func executeToolCall(call ai.ToolCall) string {
	var args map[string]interface{}
	if err := json.Unmarshal([]byte(call.Function.Arguments), &args); err != nil {
		return fmt.Sprintf("参数解析失败: %v", err)
	}

	switch call.Function.Name {
	case "get_weather":
		city, _ := args["city"].(string)
		unit, _ := args["unit"].(string)
		if unit == "" {
			unit = "celsius"
		}
		return getWeather(city, unit)

	case "calculate":
		expr, _ := args["expression"].(string)
		return calculate(expr)

	case "search_database":
		query, _ := args["query"].(string)
		limit := 5
		if l, ok := args["limit"].(float64); ok {
			limit = int(l)
		}
		return searchDatabase(query, limit)

	default:
		return fmt.Sprintf("未知的工具: %s", call.Function.Name)
	}
}

func main() {
	client := ai.NewOpenAIClient(os.Getenv("OPENAI_API_KEY"))
	ctx := context.Background()

	messages := []ai.Message{
		{Role: "system", Content: "你是一个智能助手，可以查天气、做计算和查询数据库。请根据用户的需要调用相应的工具。"},
		{Role: "user", Content: "帮我查一下北京和上海今天的天气怎么样？另外帮我算一下 sqrt(144) 是多少。"},
	}

	// 第一轮：发送用户消息，获取 AI 的工具调用决策
	fmt.Println("=== 第一轮：AI 决定调用哪些工具 ===")
	resp, err := client.Chat(ctx, &ai.ChatRequest{
		Messages:    messages,
		Tools:       tools,
		Temperature: 0,
	})
	if err != nil {
		log.Fatalf("调用失败: %v", err)
	}

	assistantMessage := resp.Choices[0].Message
	messages = append(messages, assistantMessage)

	// 检查是否有工具调用
	if len(assistantMessage.ToolCalls) == 0 {
		fmt.Println("AI 直接回答:", assistantMessage.Content)
		return
	}

	fmt.Printf("AI 决定调用 %d 个工具:\n", len(assistantMessage.ToolCalls))
	for _, tc := range assistantMessage.ToolCalls {
		fmt.Printf("  🔧 %s(%s)\n", tc.Function.Name, tc.Function.Arguments)
	}

	// 第二轮：执行所有工具调用，把结果返回给 AI
	fmt.Println("\n=== 第二轮：执行工具并返回结果 ===")
	for _, tc := range assistantMessage.ToolCalls {
		result := executeToolCall(tc)
		fmt.Printf("  📋 %s 返回: %s\n", tc.Function.Name, result)

		messages = append(messages, ai.Message{
			Role:       "tool",
			Content:    result,
			ToolCallID: tc.ID,
		})
	}

	// 第三轮：AI 根据工具结果生成最终回答
	fmt.Println("\n=== 第三轮：AI 生成最终回答 ===")
	finalResp, err := client.Chat(ctx, &ai.ChatRequest{
		Messages:    messages,
		Tools:       tools,
		Temperature: 0.7,
	})
	if err != nil {
		log.Fatalf("调用失败: %v", err)
	}

	fmt.Println(finalResp.Choices[0].Message.Content)
}

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
)

// EmbeddingRequest 是嵌入请求
type EmbeddingRequest struct {
	Model string   `json:"model"`
	Input []string `json:"input"`
}

// EmbeddingResponse 是嵌入响应
type EmbeddingResponse struct {
	Data  []EmbeddingData `json:"data"`
	Usage Usage           `json:"usage"`
}

// EmbeddingData 包含单个文本的向量
type EmbeddingData struct {
	Embedding []float64 `json:"embedding"`
	Index     int       `json:"index"`
}

// GetEmbeddings 获取文本的向量嵌入
func (c *Client) GetEmbeddings(ctx context.Context, texts []string) (*EmbeddingResponse, error) {
	req := EmbeddingRequest{
		Model: "text-embedding-3-small",
		Input: texts,
	}

	body, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.baseURL+"/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("API 错误: %s", string(respBody))
	}

	var embResp EmbeddingResponse
	if err := json.NewDecoder(resp.Body).Decode(&embResp); err != nil {
		return nil, err
	}

	return &embResp, nil
}

// CosineSimilarity 计算两个向量的余弦相似度
func CosineSimilarity(a, b []float64) float64 {
	if len(a) != len(b) {
		return 0
	}

	var dotProduct, normA, normB float64
	for i := range a {
		dotProduct += a[i] * b[i]
		normA += a[i] * a[i]
		normB += b[i] * b[i]
	}

	if normA == 0 || normB == 0 {
		return 0
	}

	return dotProduct / (math.Sqrt(normA) * math.Sqrt(normB))
}

// EuclideanDistance 计算两个向量的欧几里得距离
func EuclideanDistance(a, b []float64) float64 {
	if len(a) != len(b) {
		return math.MaxFloat64
	}

	var sum float64
	for i := range a {
		diff := a[i] - b[i]
		sum += diff * diff
	}
	return math.Sqrt(sum)
}

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"sort"

	"example.com/aiapp/ai"
)

// Document 表示一个可搜索的文档
type Document struct {
	ID        string
	Content   string
	Metadata  map[string]string
	Embedding []float64
}

// VectorStore 简单的内存向量存储
type VectorStore struct {
	documents []Document
}

func NewVectorStore() *VectorStore {
	return &VectorStore{}
}

// SearchResult 搜索结果
type SearchResult struct {
	Document Document
	Score    float64
}

// AddDocuments 向存储中添加文档
func (vs *VectorStore) AddDocuments(docs []Document) {
	vs.documents = append(vs.documents, docs...)
}

// Search 根据查询向量搜索相似文档
func (vs *VectorStore) Search(queryEmbedding []float64, topK int) []SearchResult {
	results := make([]SearchResult, len(vs.documents))

	for i, doc := range vs.documents {
		results[i] = SearchResult{
			Document: doc,
			Score:    ai.CosineSimilarity(queryEmbedding, doc.Embedding),
		}
	}

	// 按相似度降序排序
	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})

	if topK < len(results) {
		results = results[:topK]
	}

	return results
}

func main() {
	client := ai.NewOpenAIClient(os.Getenv("OPENAI_API_KEY"))
	ctx := context.Background()

	// 准备知识库文档
	texts := []string{
		"Go 语言由 Google 开发，于 2009 年开源。它的设计目标是简洁、高效、并发安全。",
		"goroutine 是 Go 语言的轻量级线程，创建成本仅为几 KB 内存，可以轻松创建数十万个。",
		"channel 是 Go 语言中 goroutine 之间通信的管道，遵循 CSP（Communicating Sequential Processes）模型。",
		"Go 的垃圾回收器使用三色标记法，Go 1.19 引入了 Pacer 改进，大幅降低了 GC 暂停时间。",
		"interface 是 Go 语言实现多态的核心机制，空接口 interface{} 可以接受任何类型的值。",
		"context 包提供了在 goroutine 树中传递截止时间、取消信号和请求级别数据的方法。",
		"Go modules 是 Go 1.11 引入的依赖管理系统，取代了之前的 GOPATH 模式。",
		"Go 1.18 引入了泛型（Generics），使用类型参数和类型约束来实现类型安全的代码复用。",
	}

	// 批量获取嵌入向量
	fmt.Println("📝 正在生成文档嵌入向量...")
	embResp, err := client.GetEmbeddings(ctx, texts)
	if err != nil {
		log.Fatalf("获取嵌入失败: %v", err)
	}

	// 构建文档存储
	store := NewVectorStore()
	docs := make([]Document, len(texts))
	for i, text := range texts {
		docs[i] = Document{
			ID:        fmt.Sprintf("doc-%d", i),
			Content:   text,
			Embedding: embResp.Data[i].Embedding,
			Metadata:  map[string]string{"source": "go-faq"},
		}
	}
	store.AddDocuments(docs)
	fmt.Printf("✅ 已索引 %d 个文档（维度: %d）\n", len(docs), len(docs[0].Embedding))

	// 语义搜索
	queries := []string{
		"Go 语言的并发机制是什么？",
		"如何处理 Go 项目的依赖？",
		"Go 的性能优化有哪些手段？",
	}

	for _, query := range queries {
		fmt.Printf("\n🔍 查询: %s\n", query)

		queryEmb, err := client.GetEmbeddings(ctx, []string{query})
		if err != nil {
			log.Fatalf("获取查询嵌入失败: %v", err)
		}

		results := store.Search(queryEmb.Data[0].Embedding, 3)
		for i, r := range results {
			fmt.Printf("  [%d] 相似度: %.4f | %s\n", i+1, r.Score, r.Document.Content[:60])
		}
	}
}

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"

	"example.com/aiapp/ai"
)

// RAGSystem 检索增强生成系统
type RAGSystem struct {
	client      *ai.Client
	vectorStore *VectorStore
	chunkSize   int
	overlap     int
}

// NewRAGSystem 创建 RAG 系统实例
func NewRAGSystem(client *ai.Client, store *VectorStore) *RAGSystem {
	return &RAGSystem{
		client:      client,
		vectorStore: store,
		chunkSize:   500,  // 每个文本块的最大字符数
		overlap:     100,  // 文本块之间的重叠字符数
	}
}

// ChunkText 将长文本切分成小块
func (r *RAGSystem) ChunkText(text string, source string) []Document {
	var chunks []Document
	words := strings.Split(text, "")
	
	for i := 0; i < len(words); i += r.chunkSize - r.overlap {
		end := i + r.chunkSize
		if end > len(words) {
			end = len(words)
		}

		chunk := strings.Join(words[i:end], "")
		if len(chunk) < 50 {
			continue // 跳过过短的文本块
		}

		chunks = append(chunks, Document{
			ID:      fmt.Sprintf("%s-chunk-%d", source, i),
			Content: chunk,
			Metadata: map[string]string{
				"source":     source,
				"chunk_idx":  fmt.Sprintf("%d", i),
				"total_len":  fmt.Sprintf("%d", len(words)),
			},
		})

		if end == len(words) {
			break
		}
	}

	return chunks
}

// IngestDocument 处理并索引一个文档
func (r *RAGSystem) IngestDocument(ctx context.Context, text, source string) error {
	// 1. 文本切分
	chunks := r.ChunkText(text, source)
	fmt.Printf("  📄 '%s' 切分为 %d 个文本块\n", source, len(chunks))

	// 2. 批量生成嵌入向量
	texts := make([]string, len(chunks))
	for i, chunk := range chunks {
		texts[i] = chunk.Content
	}

	embResp, err := r.client.GetEmbeddings(ctx, texts)
	if err != nil {
		return fmt.Errorf("生成嵌入失败: %w", err)
	}

	// 3. 关联向量和文档块
	for i := range chunks {
		chunks[i].Embedding = embResp.Data[i].Embedding
	}

	// 4. 添加到向量存储
	r.vectorStore.AddDocuments(chunks)
	return nil
}

// Query 执行 RAG 查询
func (r *RAGSystem) Query(ctx context.Context, question string, topK int) (string, error) {
	fmt.Printf("\n🧠 RAG 查询: %s\n", question)

	// 1. 将问题转换为向量
	queryEmb, err := r.client.GetEmbeddings(ctx, []string{question})
	if err != nil {
		return "", fmt.Errorf("生成查询嵌入失败: %w", err)
	}

	// 2. 检索相关文档
	results := r.vectorStore.Search(queryEmb.Data[0].Embedding, topK)
	fmt.Printf("  📚 检索到 %d 个相关文档块:\n", len(results))
	for i, r := range results {
		fmt.Printf("    [%d] %.4f | %s...\n", i+1, r.Score,
			truncate(r.Document.Content, 80))
	}

	// 3. 构建带有上下文的 prompt
	var contextBuilder strings.Builder
	contextBuilder.WriteString("以下是从知识库中检索到的相关信息：\n\n")
	for i, result := range results {
		if result.Score < 0.3 {
			continue // 过滤掉相关度太低的结果
		}
		contextBuilder.WriteString(fmt.Sprintf("--- 文档 %d (相关度: %.2f) ---\n", i+1, result.Score))
		contextBuilder.WriteString(result.Document.Content)
		contextBuilder.WriteString("\n\n")
	}

	systemPrompt := `你是一个专业的技术问答助手。请根据提供的上下文信息回答用户的问题。
如果上下文信息不足以回答问题，请明确说明哪些信息来自知识库，哪些来自你的通用知识。
回答要准确、有深度，并且易于理解。`

	userPrompt := fmt.Sprintf("上下文信息：\n%s\n\n用户问题：%s", contextBuilder.String(), question)

	// 4. 调用 LLM 生成回答
	resp, err := r.client.Chat(ctx, &ai.ChatRequest{
		Messages: []ai.Message{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userPrompt},
		},
		Temperature: 0.3, // RAG 场景通常使用较低的温度以提高准确性
		MaxTokens:   2000,
	})
	if err != nil {
		return "", fmt.Errorf("LLM 调用失败: %w", err)
	}

	return resp.Choices[0].Message.Content, nil
}

func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}

func main() {
	client := ai.NewOpenAIClient(os.Getenv("OPENAI_API_KEY"))
	ctx := context.Background()

	// 创建 RAG 系统
	store := NewVectorStore()
	rag := NewRAGSystem(client, store)

	// 导入知识库文档
	documents := map[string]string{
		"goroutine-basics": `
			goroutine 是 Go 语言中最核心的并发原语。每个 Go 程序至少有一个 goroutine（main goroutine）。
			创建 goroutine 非常简单，只需在函数调用前加上 go 关键字。goroutine 的栈是动态伸缩的，
			初始大小仅为 2KB（Go 1.4 之后），可以根据需要自动增长或缩小。
			Go 的调度器使用 GMP 模型：G 代表 goroutine，M 代表操作系统线程，P 代表处理器（逻辑 CPU）。
			默认情况下，P 的数量等于 CPU 核心数，可以通过 GOMAXPROCS 环境变量或 runtime.GOMAXPROCS() 函数调整。
			goroutine 之间的通信主要通过 channel 实现，遵循"不要通过共享内存来通信，而要通过通信来共享内存"的哲学。
		`,
		"channel-patterns": `
			Go 的 channel 有多种使用模式。无缓冲 channel（make(chan T)）要求发送方和接收方同时就绪。
			有缓冲 channel（make(chan T, N)）允许发送方在缓冲区未满时不阻塞。
			常见的模式包括：fan-in（多路复用）、fan-out（扇出分发）、pipeline（管道处理）和 worker pool（工作池）。
			select 语句用于在多个 channel 操作之间进行选择，类似于多路复用器。
			使用 close() 关闭 channel 后，接收方可以通过 v, ok := <-ch 的形式检测 channel 是否已关闭。
			注意：只能由发送方关闭 channel，向已关闭的 channel 发送数据会触发 panic。
		`,
		"context-package": `
			context 包是 Go 并发编程中不可或缺的工具。它主要解决三个问题：超时控制、取消传播和请求级别数据传递。
			context.Background() 创建根上下文，context.TODO() 用于还不确定使用哪个上下文时。
			context.WithCancel 创建可取消的上下文，调用 cancel() 后所有子上下文都会被取消。
			context.WithTimeout 和 context.WithDeadline 创建带超时的上下文，超时后自动取消。
			context.WithValue 用于在上下文中传递键值对数据，但不应滥用——它不是用来传递函数参数的。
			最佳实践：context 应该作为函数的第一个参数传递，且不应该存储在结构体中。
		`,
	}

	fmt.Println("📚 正在导入知识库...")
	for source, text := range documents {
		if err := rag.IngestDocument(ctx, text, source); err != nil {
			log.Fatalf("导入文档 '%s' 失败: %v", source, err)
		}
	}
	fmt.Println("✅ 知识库导入完成！")

	// 测试 RAG 查询
	questions := []string{
		"goroutine 的 GMP 调度模型是怎么工作的？",
		"如何优雅地关闭 channel？有什么注意事项？",
		"context.WithTimeout 和 context.WithCancel 有什么区别？",
	}

	for _, q := range questions {
		answer, err := rag.Query(ctx, q, 2)
		if err != nil {
			log.Printf("查询失败: %v", err)
			continue
		}
		fmt.Printf("\n💬 回答:\n%s\n", answer)
		fmt.Println(strings.Repeat("═", 60))
	}
}