mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-19 14:11:04 +08:00
995e405e8c
…end) ### Summary Support pipeline DSL modification through dataset configuration (backend) Key modification: knowledgebase.parser_config --------- Co-authored-by: yzc <yuzhichang@gmail.com>
105 lines
3.3 KiB
Go
105 lines
3.3 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package task
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/entity/models"
|
|
componentpkg "ragflow/internal/ingestion/component"
|
|
"ragflow/internal/service"
|
|
)
|
|
|
|
type embedder struct {
|
|
model *models.EmbeddingModel
|
|
}
|
|
|
|
func (e *embedder) MaxTokens() int {
|
|
if e == nil || e.model == nil {
|
|
return 0
|
|
}
|
|
return e.model.MaxTokens
|
|
}
|
|
|
|
func (e *embedder) Encode(texts []string) ([]componentpkg.EmbeddingResult, error) {
|
|
if e.model.ModelDriver == nil {
|
|
return nil, fmt.Errorf("embedder: embedding model driver is nil for model %v", e.model.ModelName)
|
|
}
|
|
config := &models.EmbeddingConfig{Dimension: 0}
|
|
embeds, err := e.model.ModelDriver.Embed(e.model.ModelName, texts, e.model.APIConfig, config)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
vecs := make([]componentpkg.EmbeddingResult, len(embeds))
|
|
for i, v := range embeds {
|
|
vecs[i] = componentpkg.EmbeddingResult{Vector: v.Embedding, TokenCount: v.TokenCount}
|
|
}
|
|
return vecs, nil
|
|
}
|
|
|
|
// newEmbedderResolver builds the production embedder resolver used by the
|
|
// Tokenizer component. It always resolves the embedder from the dataset's
|
|
// configured embd_id (looked up by kbID). If the dataset has no embd_id
|
|
// configured, it returns nil (no embedding). Kept as a constructor over
|
|
// injectable deps so the resolution logic stays unit-testable without a live
|
|
// model provider / DB.
|
|
func newEmbedderResolver(
|
|
getKBEmbdID func(kbID string) (string, error),
|
|
getEmbeddingModel func(tenantID, embdID string) (*models.EmbeddingModel, error),
|
|
) componentpkg.EmbedderResolver {
|
|
return func(tenantID, kbID, _ string) (componentpkg.Embedder, error) {
|
|
embdID, err := getKBEmbdID(kbID)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("embedder: resolve kb embd_id for kb_id=%s: %w", kbID, err)
|
|
}
|
|
embdID = strings.TrimSpace(embdID)
|
|
if embdID == "" {
|
|
return nil, nil
|
|
}
|
|
model, err := getEmbeddingModel(tenantID, embdID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if model == nil {
|
|
return nil, fmt.Errorf("embedder: resolved embedding model is nil for embd_id=%s", embdID)
|
|
}
|
|
return &embedder{model: model}, nil
|
|
}
|
|
}
|
|
|
|
// init wires the production embedder resolver into the component package. The
|
|
// component package must not import internal/service (dependency direction),
|
|
// so the concrete resolver is injected here - the task package is the
|
|
// composition root for ingestion runs.
|
|
func init() {
|
|
componentpkg.DefaultEmbedderResolver = newEmbedderResolver(
|
|
func(kbID string) (string, error) {
|
|
kb, err := dao.NewKnowledgebaseDAO().GetByID(kbID)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if kb == nil {
|
|
return "", nil
|
|
}
|
|
return kb.EmbdID, nil
|
|
},
|
|
service.NewModelProviderService().GetEmbeddingModel,
|
|
)
|
|
}
|