mẫu pgvector retriever

Bạn có thể sử dụng PostgreSQL và pgvector làm cách triển khai trình truy xuất. Sử dụng các ví dụ sau làm điểm xuất phát và sửa đổi để hoạt động với cơ sở dữ liệu của bạn giản đồ.

Chúng tôi sử dụng database/sql để kết nối với máy chủ Postgres nhưng bạn có thể sử dụng thư viện ứng dụng khác tùy chọn.

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)
}

Và dưới đây là cách sử dụng chú chó trong một dòng dữ liệu:

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
})