init push

backend/store/rag/rag.go

@@ -0,0 +1,74 @@
+package rag
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/cloudwego/eino/schema"
+	"github.com/google/wire"
+
+	"github.com/chaitin/panda-wiki/config"
+	"github.com/chaitin/panda-wiki/domain"
+	"github.com/chaitin/panda-wiki/log"
+)
+
+type QueryRecordsRequest struct {
+	DatasetID           string
+	Query               string
+	GroupIDs            []int
+	Tags                []string
+	SimilarityThreshold float64
+	HistoryMsgs         []*schema.Message
+	MaxChunksPerDoc     int
+}
+
+type UpsertRecordsRequest struct {
+	ID        string
+	DatasetID string
+	DocID     string
+	Title     string
+	Content   string
+	GroupIDs  []int
+	Tags      []string
+}
+
+type DocumentMetadata struct {
+	GroupIDs []int `json:"group_ids"`
+}
+
+type Document struct {
+	ID          string           `json:"id"`
+	Name        string           `json:"name"`
+	DatasetID   string           `json:"dataset_id"`
+	Status      string           `json:"status"`
+	ProgressMsg string           `json:"progress_msg"`
+	MetaData    DocumentMetadata `json:"meta_data"`
+	Tags        []string         `json:"tags"`
+}
+
+type RAGService interface {
+	CreateKnowledgeBase(ctx context.Context) (string, error)
+	UpsertRecords(ctx context.Context, req *UpsertRecordsRequest) (string, error)
+	QueryRecords(ctx context.Context, req *QueryRecordsRequest) (string, []*domain.NodeContentChunk, error)
+	DeleteRecords(ctx context.Context, datasetID string, docIDs []string) error
+	DeleteKnowledgeBase(ctx context.Context, datasetID string) error
+	UpdateDocumentGroupIDs(ctx context.Context, datasetID string, docID string, groupIds []int) error
+	ListDocuments(ctx context.Context, datasetID string, documentIDs []string) ([]Document, error)
+
+	GetModelList(ctx context.Context) ([]*domain.Model, error)
+	AddModel(ctx context.Context, model *domain.Model) (string, error)
+	UpdateModel(ctx context.Context, model *domain.Model) error
+	UpsertModel(ctx context.Context, model *domain.Model) error
+	DeleteModel(ctx context.Context, model *domain.Model) error
+}
+
+func NewRAGService(config *config.Config, logger *log.Logger) (RAGService, error) {
+	switch config.RAG.Provider {
+	case "ct":
+		return NewCTRAG(config, logger)
+	default:
+		return nil, fmt.Errorf("unsupported vector provider: %s", config.RAG.Provider)
+	}
+}
+
+var ProviderSet = wire.NewSet(NewRAGService)