Шаблон ретривера pgvector

Вы можете использовать PostgreSQL и pgvector в качестве реализации средства извлечения. Используйте следующие примеры в качестве отправной точки и измените их для работы со схемой вашей базы данных.

Мы используем базу данных/sql для подключения к серверу Postgres, но вы можете использовать другую клиентскую библиотеку по вашему выбору.

func defineRetriever(db *sql.DB, embedder *ai.Embedder) *ai.Retriever {
  f := func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) {
      eres, err := embedder.Embed(ctx, &ai.EmbedRequest{Documents: []*ai.Document{req.Document}})
      if err != nil {
          return nil, err
      }
      rows, err := db.QueryContext(ctx, `
          SELECT episode_id, season_number, chunk as content
          FROM embeddings
          WHERE show_id = $1
          ORDER BY embedding <#> $2
          LIMIT 2`,
          req.Options, pgv.NewVector(eres.Embeddings[0].Embedding))
      if err != nil {
          return nil, err
      }
      defer rows.Close()

      res := &ai.RetrieverResponse{}
      for rows.Next() {
          var eid, sn int
          var content string
          if err := rows.Scan(&eid, &sn, &content); err != nil {
              return nil, err
          }
          meta := map[string]any{
              "episode_id":    eid,
              "season_number": sn,
          }
          doc := &ai.Document{
              Content:  []*ai.Part{ai.NewTextPart(content)},
              Metadata: meta,
          }
          res.Documents = append(res.Documents, doc)
      }
      if err := rows.Err(); err != nil {
          return nil, err
      }
      return res, nil
  }
  return ai.DefineRetriever(provider, "shows", f)
}

А вот как использовать ретривер в потоке:

retriever := defineRetriever(db, embedder)

type input struct {
  Question string
  Show     string
}

genkit.DefineFlow("askQuestion", func(ctx context.Context, in input) (string, error) {
  res, err := retriever.Retrieve(ctx, &ai.RetrieverRequest{
      Document: &ai.Document{Content: []*ai.Part{ai.NewTextPart(in.Question)}},
      Options:  in.Show,
  })
  if err != nil {
      return "", err
  }
  for _, doc := range res.Documents {
      fmt.Printf("%+v %q\n", doc.Metadata, doc.Content[0].Text)
  }
  // Use documents in RAG prompts.
  return "", nil
})