使用golang实现了一整套 rag 功能

基于eino+es实现了一套 rag 功能

包括 文件内容解析、chunk、embedding、retriever

GitHub - wangle201210/go-rag: 基于eino实现知识库的rag

点个⭐️一起学习

go-rag

基于eino实现知识库的rag

存储层

  • es8存储向量相关数据

功能列表

  • md、pdf、html 文档解析
  • 网页解析
  • 文档检索

未来计划

  • 使用mysql存储chunk和文档的映射关系,目前放在es的ext字段

使用

安装依赖

go get github.com/wangle201210/go-rag@latest

安装es8

docker run -d --name elasticsearch \
  -e "discovery.type=single-node" \
  -e "ES_JAVA_OPTS=-Xms512m -Xmx512m" \
  -e "xpack.security.enabled=false" \
  -p 9200:9200 \
  -p 9300:9300 \
  elasticsearch:8.18.0

配置环境变量(也可以在文件中直接配置)

export OPENAI_API_KEY=YOUR_OPENAI_API_KEY
export OPENAI_BASE_URL=YOUR_OPENAI_BASE_URL

加载各种数据源的数据,并将其向量化后存储进向量数据库。

// Index
// uri: 文档地址,可以是文件路径(pdf,html,md等),也可以是网址
// ids: 文档id
func Index(uri string) (ids []string, err error) {
    buildIndex, err := indexer.BuildIndexer(context.Background())
    if err != nil {
        return
    }
    s := document.Source{
        URI: uri,
    }
  // 这个buildIndex可以复用
    ids, err = buildIndex.Invoke(context.Background(), s)
    if err != nil {
        return
    }
    return
}

检索

// Retrieve
// input: 检索关键词
// score: 0-2, 0 完全相反,1 毫不相干,2 完全相同
func Retrieve(input string, score float64) (msg []*schema.Document, err error) {
  r, err := retriever.BuildRetriever(context.Background())
  if err != nil {
    return
  }
  // 这个r可以复用
  msg, err = r.Invoke(context.Background(), input,
    compose.WithRetrieverOption(
      er.WithScoreThreshold(score),
      er.WithTopK(5),
    ),
  )
  if err != nil {
    return
  }
  return
}

详情可以参照test文件

本作品采用《CC 协议》,转载必须注明作者和本文链接
讨论数量: 1

讨论应以学习和精进为目的。请勿发布不友善或者负能量的内容,与人为善,比聪明更重要!