ben/coze-studio
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(knowledge): Move the all dependent components to app infra (#795)

Browse Source
This commit is contained in:
Ryo
2025-08-18 18:18:39 +08:00
committed by GitHub
parent 23a468c72c
commit f940edf585
9 changed files with 198 additions and 186 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
backend/application/base/appinfra/app_infra.go
View File

@@ -18,9 +18,11 @@ package appinfra
import (
"context"
"encoding/json"
"fmt"
"net/http"
"os"
"path/filepath"
"strconv"
"strings"
"time"
@@ -31,26 +33,32 @@ import (
"github.com/cloudwego/eino-ext/components/embedding/gemini"
"github.com/cloudwego/eino-ext/components/embedding/ollama"
"github.com/cloudwego/eino-ext/components/embedding/openai"
"github.com/cloudwego/eino/components/prompt"
"github.com/cloudwego/eino/schema"
"github.com/milvus-io/milvus/client/v2/milvusclient"
"github.com/volcengine/volc-sdk-golang/service/visual"
"github.com/coze-dev/coze-studio/backend/application/internal"
"github.com/coze-dev/coze-studio/backend/infra/contract/cache"
"github.com/coze-dev/coze-studio/backend/infra/contract/chatmodel"
"github.com/coze-dev/coze-studio/backend/infra/contract/coderunner"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/nl2sql"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/ocr"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/parser"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/rerank"
"github.com/coze-dev/coze-studio/backend/infra/contract/document/searchstore"
"github.com/coze-dev/coze-studio/backend/infra/contract/embedding"
"github.com/coze-dev/coze-studio/backend/infra/contract/imagex"
"github.com/coze-dev/coze-studio/backend/infra/contract/messages2query"
"github.com/coze-dev/coze-studio/backend/infra/contract/modelmgr"
"github.com/coze-dev/coze-studio/backend/infra/impl/cache/redis"
"github.com/coze-dev/coze-studio/backend/infra/impl/coderunner/direct"
"github.com/coze-dev/coze-studio/backend/infra/impl/coderunner/sandbox"
builtinNL2SQL "github.com/coze-dev/coze-studio/backend/infra/impl/document/nl2sql/builtin"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/ocr/ppocr"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/ocr/veocr"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/parser/builtin"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/parser/ppstructure"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/rerank/rrf"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/searchstore/elasticsearch"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/searchstore/milvus"
"github.com/coze-dev/coze-studio/backend/infra/impl/document/searchstore/vikingdb"
@@ -61,6 +69,7 @@ import (
"github.com/coze-dev/coze-studio/backend/infra/impl/eventbus"
"github.com/coze-dev/coze-studio/backend/infra/impl/idgen"
"github.com/coze-dev/coze-studio/backend/infra/impl/imagex/veimagex"
builtinM2Q "github.com/coze-dev/coze-studio/backend/infra/impl/messages2query/builtin"
"github.com/coze-dev/coze-studio/backend/infra/impl/mysql"
"github.com/coze-dev/coze-studio/backend/infra/impl/storage"
"github.com/coze-dev/coze-studio/backend/pkg/lang/conv"
@@ -70,19 +79,24 @@ import (
)
type AppDependencies struct {
DB *gorm.DB
CacheCli cache.Cmdable
IDGenSVC idgen.IDGenerator
ESClient es.Client
ImageXClient imagex.ImageX
TOSClient storage.Storage
ResourceEventProducer eventbus.Producer
AppEventProducer eventbus.Producer
ModelMgr modelmgr.Manager
CodeRunner coderunner.Runner
OCR ocr.OCR
ParserManager parser.Manager
SearchStoreManagers []searchstore.Manager
DB *gorm.DB
CacheCli cache.Cmdable
IDGenSVC idgen.IDGenerator
ESClient es.Client
ImageXClient imagex.ImageX
TOSClient storage.Storage
ResourceEventProducer eventbus.Producer
AppEventProducer eventbus.Producer
KnowledgeEventProducer eventbus.Producer
ModelMgr modelmgr.Manager
CodeRunner coderunner.Runner
OCR ocr.OCR
ParserManager parser.Manager
SearchStoreManagers []searchstore.Manager
Reranker rerank.Reranker
Rewriter messages2query.MessagesToQuery
NL2SQL nl2sql.NL2SQL
WorkflowBuildInChatModel chatmodel.BaseChatModel
}
func Init(ctx context.Context) (*AppDependencies, error) {
@@ -126,6 +140,23 @@ func Init(ctx context.Context) (*AppDependencies, error) {
return nil, fmt.Errorf("init app event producer failed, err=%w", err)
}
deps.KnowledgeEventProducer, err = initKnowledgeEventBusProducer()
if err != nil {
return nil, fmt.Errorf("init knowledge event bus producer failed, err=%w", err)
}
deps.Reranker = rrf.NewRRFReranker(0)
deps.Rewriter, err = initRewriter(ctx)
if err != nil {
return nil, fmt.Errorf("init rewriter failed, err=%w", err)
}
deps.NL2SQL, err = initNL2SQL(ctx)
if err != nil {
return nil, fmt.Errorf("init nl2sql failed, err=%w", err)
}
deps.ModelMgr, err = initModelMgr()
if err != nil {
return nil, fmt.Errorf("init model manager failed, err=%w", err)
@@ -135,11 +166,21 @@ func Init(ctx context.Context) (*AppDependencies, error) {
deps.OCR = initOCR()
imageAnnotationModel, _, err := internal.GetBuiltinChatModel(ctx, "IA_")
imageAnnotationModel, _, err := getBuiltinChatModel(ctx, "IA_")
if err != nil {
return nil, fmt.Errorf("get builtin chat model failed, err=%w", err)
}
var ok bool
deps.WorkflowBuildInChatModel, ok, err = getBuiltinChatModel(ctx, "WKR_")
if err != nil {
return nil, fmt.Errorf("get workflow builtin chat model failed, err=%w", err)
}
if !ok {
logs.CtxWarnf(ctx, "workflow builtin chat model for knowledge recall not configured")
}
deps.ParserManager, err = initParserManager(deps.TOSClient, deps.OCR, imageAnnotationModel)
if err != nil {
return nil, fmt.Errorf("init parser manager failed, err=%w", err)
@@ -166,6 +207,71 @@ func initSearchStoreManagers(ctx context.Context, es es.Client) ([]searchstore.M
return []searchstore.Manager{esSearchstoreManager, mgr}, nil
}
func initRewriter(ctx context.Context) (messages2query.MessagesToQuery, error) {
rewriterChatModel, _, err := getBuiltinChatModel(ctx, "M2Q_")
if err != nil {
return nil, err
}
filePath := filepath.Join(getWorkingDirectory(), "resources/conf/prompt/messages_to_query_template_jinja2.json")
rewriterTemplate, err := readJinja2PromptTemplate(filePath)
if err != nil {
return nil, err
}
rewriter, err := builtinM2Q.NewMessagesToQuery(ctx, rewriterChatModel, rewriterTemplate)
if err != nil {
return nil, err
}
return rewriter, nil
}
func getWorkingDirectory() string {
root, err := os.Getwd()
if err != nil {
logs.Warnf("[InitConfig] Failed to get current working directory: %v", err)
root = os.Getenv("PWD")
}
return root
}
func readJinja2PromptTemplate(jsonFilePath string) (prompt.ChatTemplate, error) {
b, err := os.ReadFile(jsonFilePath)
if err != nil {
return nil, err
}
var m2qMessages []*schema.Message
if err = json.Unmarshal(b, &m2qMessages); err != nil {
return nil, err
}
tpl := make([]schema.MessagesTemplate, len(m2qMessages))
for i := range m2qMessages {
tpl[i] = m2qMessages[i]
}
return prompt.FromMessages(schema.Jinja2, tpl...), nil
}
func initNL2SQL(ctx context.Context) (nl2sql.NL2SQL, error) {
n2sChatModel, _, err := getBuiltinChatModel(ctx, "NL2SQL_")
if err != nil {
return nil, err
}
filePath := filepath.Join(getWorkingDirectory(), "resources/conf/prompt/nl2sql_template_jinja2.json")
n2sTemplate, err := readJinja2PromptTemplate(filePath)
if err != nil {
return nil, err
}
n2s, err := builtinNL2SQL.NewNL2SQL(ctx, n2sChatModel, n2sTemplate)
if err != nil {
return nil, err
}
return n2s, nil
}
func initImageX(ctx context.Context) (imagex.ImageX, error) {
uploadComponentType := os.Getenv(consts.FileUploadComponentType)
if uploadComponentType != consts.FileUploadComponentTypeImagex {
@@ -206,6 +312,17 @@ func initAppEventProducer() (eventbus.Producer, error) {
return appEventProducer, nil
}
func initKnowledgeEventBusProducer() (eventbus.Producer, error) {
nameServer := os.Getenv(consts.MQServer)
knowledgeProducer, err := eventbus.NewProducer(nameServer, consts.RMQTopicKnowledge, consts.RMQConsumeGroupKnowledge, 2)
if err != nil {
return nil, fmt.Errorf("init knowledge producer failed, err=%w", err)
}
return knowledgeProducer, nil
}
func initCodeRunner() coderunner.Runner {
switch typ := os.Getenv(consts.CodeRunnerType); typ {
case "sandbox":

109
backend/application/base/appinfra/builtin_chat_model.go Normal file
View File

@@ -0,0 +1,109 @@
/*
* Copyright 2025 coze-dev Authors
*
* 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 appinfra
import (
"context"
"fmt"
"os"
"strconv"
ao "github.com/cloudwego/eino-ext/components/model/ark"
"github.com/cloudwego/eino-ext/components/model/deepseek"
"github.com/cloudwego/eino-ext/components/model/gemini"
"github.com/cloudwego/eino-ext/components/model/ollama"
mo "github.com/cloudwego/eino-ext/components/model/openai"
"github.com/cloudwego/eino-ext/components/model/qwen"
"google.golang.org/genai"
"github.com/coze-dev/coze-studio/backend/infra/contract/chatmodel"
)
func getBuiltinChatModel(ctx context.Context, envPrefix string) (bcm chatmodel.BaseChatModel, configured bool, err error) {
getEnv := func(key string) string {
if val := os.Getenv(envPrefix + key); val != "" {
return val
}
return os.Getenv(key)
}
switch getEnv("BUILTIN_CM_TYPE") {
case "openai":
byAzure, _ := strconv.ParseBool(getEnv("BUILTIN_CM_OPENAI_BY_AZURE"))
bcm, err = mo.NewChatModel(ctx, &mo.ChatModelConfig{
APIKey: getEnv("BUILTIN_CM_OPENAI_API_KEY"),
ByAzure: byAzure,
BaseURL: getEnv("BUILTIN_CM_OPENAI_BASE_URL"),
Model: getEnv("BUILTIN_CM_OPENAI_MODEL"),
})
case "ark":
bcm, err = ao.NewChatModel(ctx, &ao.ChatModelConfig{
APIKey: getEnv("BUILTIN_CM_ARK_API_KEY"),
Model: getEnv("BUILTIN_CM_ARK_MODEL"),
BaseURL: getEnv("BUILTIN_CM_ARK_BASE_URL"),
})
case "deepseek":
bcm, err = deepseek.NewChatModel(ctx, &deepseek.ChatModelConfig{
APIKey: getEnv("BUILTIN_CM_DEEPSEEK_API_KEY"),
BaseURL: getEnv("BUILTIN_CM_DEEPSEEK_BASE_URL"),
Model: getEnv("BUILTIN_CM_DEEPSEEK_MODEL"),
})
case "ollama":
bcm, err = ollama.NewChatModel(ctx, &ollama.ChatModelConfig{
BaseURL: getEnv("BUILTIN_CM_OLLAMA_BASE_URL"),
Model: getEnv("BUILTIN_CM_OLLAMA_MODEL"),
})
case "qwen":
bcm, err = qwen.NewChatModel(ctx, &qwen.ChatModelConfig{
APIKey: getEnv("BUILTIN_CM_QWEN_API_KEY"),
BaseURL: getEnv("BUILTIN_CM_QWEN_BASE_URL"),
Model: getEnv("BUILTIN_CM_QWEN_MODEL"),
})
case "gemini":
backend, convErr := strconv.ParseInt(getEnv("BUILTIN_CM_GEMINI_BACKEND"), 10, 64)
if convErr != nil {
return nil, false, convErr
}
c, clientErr := genai.NewClient(ctx, &genai.ClientConfig{
APIKey: getEnv("BUILTIN_CM_GEMINI_API_KEY"),
Backend: genai.Backend(backend),
Project: getEnv("BUILTIN_CM_GEMINI_PROJECT"),
Location: getEnv("BUILTIN_CM_GEMINI_LOCATION"),
HTTPOptions: genai.HTTPOptions{
BaseURL: getEnv("BUILTIN_CM_GEMINI_BASE_URL"),
},
})
if clientErr != nil {
return nil, false, clientErr
}
bcm, err = gemini.NewChatModel(ctx, &gemini.Config{
Client: c,
Model: getEnv("BUILTIN_CM_GEMINI_MODEL"),
})
default:
// accept builtin chat model not configured
}
if err != nil {
return nil, false, fmt.Errorf("builtin %s chat model init failed, %w", envPrefix, err)
}
if bcm != nil {
configured = true
}
return
}