核心流程--向量数据库VikingDB-火山引擎

文档中心
向量数据库VikingDB
Go SDK
核心流程
本页面提供一个通过 Go SDK 完成知识库文档导入与 RAG 问答的完整流程请求示例，涵盖文档上传、向量检索及大模型问答调用，注意使用时根据实际情况填充账户鉴权信息、collection 名称、文档路径和查询内容。
package main


import (
    "context"
    "encoding/json"
    "fmt"
    "os"
    "time"


    knowledge "github.com/volcengine/vikingdb-go-sdk/knowledge"
    kmodel "github.com/volcengine/vikingdb-go-sdk/knowledge/model"
)


func requireEnv(name string) (string, error) {
    v := os.Getenv(name)
    if v == "" {
        return "", fmt.Errorf("missing required env: %s", name)
    }
    return v, nil
}


func initClient() (*knowledge.Client, error) {
    ak, err := requireEnv("VOLC_AK")
    if err != nil {
        return nil, err
    }
    sk, err := requireEnv("VOLC_SK")
    if err != nil {
        return nil, err
    }


    endpoint := "https://api-knowledgebase.mlp.cn-beijing.volces.com"
    region := "cn-beijing"


    return knowledge.New(
        knowledge.AuthIAM(ak, sk),
        knowledge.WithEndpoint(endpoint),
        knowledge.WithRegion(region),
    )
}


func initCollection(c *knowledge.Client) *knowledge.CollectionClient {
    resourceID := "your-collection-resource-id"
    collectionName := "your-collection-name"
    projectName := "default"
    return c.Collection(kmodel.CollectionMeta{
        ResourceID:     resourceID,
        CollectionName: collectionName,
        ProjectName:    projectName,
    })
}


func addDocV2(ctx context.Context, kc *knowledge.CollectionClient, docID, docName, uri string) (*kmodel.AddDocResponse, error) {
    req := kmodel.AddDocV2Request{
        DocID:   &docID,
        DocName: nil,
        URI:     &uri,
    }
    if docName != "" {
        req.DocName = &docName
    }
    resp, err := kc.AddDocV2(ctx, req)
    if err != nil {
        return nil, err
    }
    return resp, nil
}


func makeMessages(sk *kmodel.SearchKnowledgeResponse, query string) []kmodel.ChatMessage {
    top := []string{}
    if sk != nil && sk.Data != nil {
        for i, item := range sk.Data.ResultList {
            if i >= 5 {
                break
            }
            title := ""
            if item.ChunkTitle != nil {
                title = *item.ChunkTitle
            }
            content := ""
            if item.Content != nil {
                content = *item.Content
            }
            top = append(top, fmt.Sprintf("【%s】\n%s", title, content))
        }
    }
    contextText := "（检索结果为空或不可用）"
    if len(top) > 0 {
        contextText = ""
        for idx, s := range top {
            if idx > 0 {
                contextText += "\n\n"
            }
            contextText += s
        }
    }
    systemPrompt := "你是一位专业的财报分析师，你需要根据「参考资料」来回答接下来的「用户问题」，这些信息在 <context></context> XML 标签之内。回答必须在参考资料范围内，尽可能简洁，无法回答时请礼貌说明并引导提供更多信息。\n\n<context>\n" + contextText + "\n</context>"
    return []kmodel.ChatMessage{
        {Role: "system", Content: systemPrompt},
        {Role: "user", Content: query},
    }
}


func runOverall() error {
    ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
    defer cancel()


    client, err := initClient()
    if err != nil {
        return err
    }
    kc := initCollection(client)
    if kc == nil {
        return fmt.Errorf("collection client is nil")
    }


    tosURI := "tos://your-bucket/your-path/your-file.pdf"
    tosDocID := "your-doc-id"
    tosResp, err := addDocV2(ctx, kc, tosDocID, "", tosURI)
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
    } else if tosResp.Data != nil && tosResp.Data.DocID != nil {
        fmt.Println("add tos doc id:", *tosResp.Data.DocID)
    }


    urlURI := "https://your-url.pdf"
    urlDocID := "your-doc-id"
    urlDocName := "your-file-name.pdf"


    urlResp, err := addDocV2(ctx, kc, urlDocID, urlDocName, urlURI)
    if err != nil {
        fmt.Fprintln(os.Stderr, err)
    } else if urlResp.Data != nil && urlResp.Data.DocID != nil {
        fmt.Println("add url doc id:", *urlResp.Data.DocID)
    }


    query := "your query"
    limit := 10
    denseWeight := 0.5


    skReq := kmodel.SearchKnowledgeRequest{
        Query:       query,
        Limit:       &limit,
        DenseWeight: &denseWeight,
    }
    skResp, err := kc.SearchKnowledge(ctx, skReq)
    if err != nil {
        return err
    }


    msgs := makeMessages(skResp, query)
    stream := false
    chatReq := kmodel.ChatCompletionRequest{
        Model:    "Doubao-1-5-pro-32k",
        Messages: msgs,
        Stream:   &stream,
    }
    chatResp, err := client.ChatCompletion(ctx, chatReq)
    if err != nil {
        return err
    }


    b, err := json.MarshalIndent(chatResp, "", "  ")
    if err != nil {
        return err
    }
    fmt.Println("chat_completion:", string(b))


    return nil
}


func main() {
    if err := runOverall(); err != nil {
        fmt.Fprintln(os.Stderr, err)
        os.Exit(1)
    }
}
最近更新时间：2026.03.20 12:23:12
这个页面对您有帮助吗？
有用
无用