Kimi API 助手的氮气加速装置 —— 以 Golang 为例实践 Context Caching 3

type Client interface {
	// RetrieveContextCacheTag GET {{ $.Client.BaseUrl }}/caching/refs/tags/{{ $.tag }}
	// Content-Type: application/json
	// Authorization: Bearer {{ $.Client.Key }}
	RetrieveContextCacheTag(ctx context.Context, tag string) (*ContextCacheTag, error) // <-- 在本次示例中，我们仅使用了 RetrieveContextCacheTag 接口
}
 
type CreateContextCacheRequest struct {
	Messages    []*Message        `json:"messages"`
	Model       string            `json:"model"`
	Tools       []*Tool           `json:"tools"`
	Name        string            `json:"name,omitempty"`
	Description string            `json:"description,omitempty"`
	Metadata    map[string]string `json:"metadata,omitempty"`
	ExpiredAt   int               `json:"expiredAt,omitempty"`
	TTL         int               `json:"ttl,omitempty"`
 
	Tags        []string          `json:"tags,omitempty"` // <-- 我们为 CreateContextCacheRequest 添加了一个额外的 tags 字段，以便于在创建 Cache 时指定 Tag
}
 
// ContextCacheTag 用于记录 Tag 与 Cache 的映射关系
type ContextCacheTag struct {
	Tag     string `json:"tag"`
	CacheID string `json:"cache_id"`
}

var documents = []string{
	"kimi-api-doc-001.pdf",
	"kimi-api-doc-002.pdf",
	"kimi-api-doc-003.pdf",
}
 
func initKnowledge(ctx context.Context, client Client, tag string) (err error) {
	// 通过 RetrieveContextCacheTag 接口，我们检查当前 tag 是否有可用的 Cache，并通过返回结果决定我们是新增 Cache、或是
	// 使用 tag 绑定的 Cache。
	_, err = client.RetrieveContextCacheTag(ctx, tag)
	if err != nil {
		// 通过 error.type 判断 tag 对应的 Cache 是否存在，在 tag 系统的设计中，一个 tag 必须与一个 cache_id 绑定，
		// 如果查询不到 tag，意味着我们还没有创建 Cache ，需要初始化 Cache 并与当前 tag 绑定。
		//
		// 我们将需要 Caching 的文件逐个通过文件抽取接口获得文件内容，然后再组装成 messages 创建 Cache，在创建 Cache 
		// 时，通过 tags 字段指定 tag。
		if parsed := ParseError(err); parsed != nil && parsed.Type == "resource_not_found_error" {
			// 逐个抽取文件内容，并组装成 messages
			messages := make([]*Message, 0, len(documents))
			for _, document := range documents {
				file, err := os.Open(document)
				if err != nil {
					return err
				}
				uploadedFile, err := client.UploadFile(ctx, &UploadFileRequest{
					File:    file,
					Purpose: "file-extract",
				})
				file.Close()
				if err != nil {
					return err
				}
				fileContent, err := client.RetrieveFileContent(ctx, uploadedFile.ID)
				if err != nil {
					return err
				}
				messages = append(messages, &Message{
					Role:    RoleSystem,
					Content: &Content{Text: string(fileContent)},
				})
			}
 
			// 使用组装好的 messages 创建缓存，并打上 tag
			_, err = client.CreateContextCache(ctx, &CreateContextCacheRequest{
				Messages: messages,
				Model:    "moonshot-v1",
				TTL:      3600,
				Tags:     []string{tag},
			})
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}
 
	// 如果当前 tag 已经绑定 Cache，我们不需要做任何额外操作，直接使用这个 tag 即可（返回 nil error）
	return nil
}

// ContextCacheOptions 依然是在上一个章节就出现过的老朋友了
type ContextCacheOptions struct {
	CacheID  string
	Tag      string
	ResetTTL int
}
 
type Content struct {
	Text  string
	Cache *ContextCacheOptions
}
 
func (c *Content) MarshalJSON() ([]byte, error) {
	if c == nil {
		return json.Marshal(nil)
	}
	if c.Cache != nil {
		var cacheOptionsBuilder strings.Builder
		if tag := c.Cache.Tag; tag != "" {
			cacheOptionsBuilder.WriteString("tag=")
			cacheOptionsBuilder.WriteString(c.Cache.Tag)
		} else {
			cacheOptionsBuilder.WriteString("cache_id=")
			cacheOptionsBuilder.WriteString(c.Cache.CacheID)
		}
		if resetTTL := c.Cache.ResetTTL; resetTTL > 0 {
			cacheOptionsBuilder.WriteString(";")
			cacheOptionsBuilder.WriteString("reset_ttl=")
			cacheOptionsBuilder.WriteString(strconv.Itoa(resetTTL))
		}
		return json.Marshal(cacheOptionsBuilder.String())
	}
	return json.Marshal(c.Text)
}
 
func chatWithCacheByMessage(
	ctx context.Context,
	client Client[moonshot],
	tag string,
) {
	if err := initKnowledge(ctx, client, tag); err != nil {
		if parsed := ParseError(err); parsed != nil {
			log.Fatalln("("+parsed.Type+")", parsed.Message)
		}
		log.Fatalln(err)
	}
 
	messages = []*Message{
		{Role: RoleCache, Content: &Content{Cache: &ContextCacheOptions{Tag: tag, ResetTTL: 3600}}},
		{Role: RoleUser, Content: &Content{Text: input}},
	}
 
	stream, err := client.CreateChatCompletionStream(ctx, &ChatCompletionStreamRequest{
		Messages: messages,
		Model:    ModelMoonshot128K,
	})
 
	if err != nil {
		if parsed := ParseError(err); parsed != nil {
			log.Fatalln("("+parsed.Type+")", parsed.Message)
		}
		log.Fatalln(err)
	}
 
	defer stream.Close()
	for chunk := range stream.C {
		fmt.Printf("%s", chunk.GetDeltaContent())
	}
}

func updateKnowledge(ctx context.Context, client Client[moonshot], tag string) (err error) {
	// 步骤一：使用小助手的 tag 查询对应的 cache_id
	cacheTag, err := client.RetrieveContextCacheTag(ctx, tag)
 
	// 额外的，我们会忽略在查询 tag 过程中产生的错误，我们简单地认为，不管遇到什么错误，
	// 我们都默认需要重新创建 Cache。
	if err == nil && cacheTag != nil {
		// 步骤二：删除 cache_id 对应的 Cache
		if err = client.DeleteContextCache(ctx, cacheTag.CacheID); err != nil {
			return err
		}
	}
 
	// 步骤三：使用新的知识创建一个新的 Cache，并为其打上与步骤一相同的 tag
	return initKnowledge(ctx, client, tag)
}