diff --git a/graphiti b/graphiti new file mode 160000 index 000000000..d2654003f --- /dev/null +++ b/graphiti @@ -0,0 +1 @@ +Subproject commit d2654003ffc11821bce73c493162a40181b23504 diff --git a/internal/application/service/extract.go b/internal/application/service/extract.go index c7fa67ce6..cd8989116 100644 --- a/internal/application/service/extract.go +++ b/internal/application/service/extract.go @@ -18,7 +18,7 @@ import ( func NewChunkExtractTask(ctx context.Context, client *asynq.Client, tenantID uint, chunkID string, modelID string) error { if strings.ToLower(os.Getenv("NEO4J_ENABLE")) != "true" { - logger.Debugf(ctx, "NOT SUPPORT RETRIEVE GRAPH") + logger.Warn(ctx, "NEO4J is not enabled, skip chunk extract task") return nil } payload, err := json.Marshal(types.ExtractChunkPayload{ diff --git a/internal/application/service/knowledge.go b/internal/application/service/knowledge.go index 490fa177b..d548053f8 100644 --- a/internal/application/service/knowledge.go +++ b/internal/application/service/knowledge.go @@ -9,7 +9,6 @@ import ( "fmt" "io" "mime/multipart" - "os" "regexp" "slices" "sort" @@ -1045,8 +1044,6 @@ func (s *knowledgeService) processChunks(ctx context.Context, return } - enableGraphRAG := os.Getenv("ENABLE_GRAPH_RAG") == "true" - // Create chunk objects from proto chunks maxSeq := 0 @@ -1187,55 +1184,6 @@ func (s *knowledgeService) processChunks(ctx context.Context, textChunks[i+1].PreChunkID = chunk.ID } } - if enableGraphRAG { - relationChunkSize := 5 - indirectRelationChunkSize := 5 - graphBuilder := NewGraphBuilder(s.config, chatModel) - err = graphBuilder.BuildGraph(ctx, textChunks) - if err != nil { - logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks build graph failed") - span.RecordError(err) - } else { - for _, chunk := range textChunks { - chunk.RelationChunks, _ = json.Marshal(graphBuilder.GetRelationChunks(chunk.ID, relationChunkSize)) - chunk.IndirectRelationChunks, _ = json.Marshal(graphBuilder.GetIndirectRelationChunks(chunk.ID, indirectRelationChunkSize)) - } - for i, entity := range graphBuilder.GetAllEntities() { - relationChunks, _ := json.Marshal(entity.ChunkIDs) - entityChunk := &types.Chunk{ - ID: entity.ID, - TenantID: knowledge.TenantID, - KnowledgeID: knowledge.ID, - KnowledgeBaseID: knowledge.KnowledgeBaseID, - Content: entity.Description, - ChunkIndex: maxSeq + i*100 + 3, - IsEnabled: true, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - ChunkType: types.ChunkTypeEntity, - RelationChunks: types.JSON(relationChunks), - } - insertChunks = append(insertChunks, entityChunk) - } - for i, relationship := range graphBuilder.GetAllRelationships() { - relationChunks, _ := json.Marshal(relationship.ChunkIDs) - relationshipChunk := &types.Chunk{ - ID: relationship.ID, - TenantID: knowledge.TenantID, - KnowledgeID: knowledge.ID, - KnowledgeBaseID: knowledge.KnowledgeBaseID, - Content: relationship.Description, - ChunkIndex: maxSeq + i*100 + 4, - IsEnabled: true, - CreatedAt: time.Now(), - UpdatedAt: time.Now(), - ChunkType: types.ChunkTypeRelationship, - RelationChunks: types.JSON(relationChunks), - } - insertChunks = append(insertChunks, relationshipChunk) - } - } - } span.AddEvent("extract summary") summary, err := s.getSummary(ctx, chatModel, knowledge, textChunks) @@ -1357,11 +1305,13 @@ func (s *knowledgeService) processChunks(ctx context.Context, logger.GetLogger(ctx).Infof("processChunks batch index successfully, with %d index", len(indexInfoList)) logger.Infof(ctx, "processChunks create relationship rag task") - for _, chunk := range textChunks { - err := NewChunkExtractTask(ctx, s.task, chunk.TenantID, chunk.ID, kb.SummaryModelID) - if err != nil { - logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks create chunk extract task failed") - span.RecordError(err) + if kb.ExtractConfig.Enabled { + for _, chunk := range textChunks { + err := NewChunkExtractTask(ctx, s.task, chunk.TenantID, chunk.ID, kb.SummaryModelID) + if err != nil { + logger.GetLogger(ctx).WithField("error", err).Errorf("processChunks create chunk extract task failed") + span.RecordError(err) + } } } diff --git a/internal/handler/initialization.go b/internal/handler/initialization.go index a334c49d7..8994c9f51 100644 --- a/internal/handler/initialization.go +++ b/internal/handler/initialization.go @@ -1241,7 +1241,7 @@ func (h *InitializationHandler) buildConfigResponse(ctx context.Context, models if kb.ExtractConfig != nil { config["nodeExtract"] = map[string]interface{}{ - "enabled": true, + "enabled": kb.ExtractConfig.Enabled, "text": kb.ExtractConfig.Text, "tags": kb.ExtractConfig.Tags, "nodes": kb.ExtractConfig.Nodes, diff --git a/internal/types/knowledgebase.go b/internal/types/knowledgebase.go index 7d076849a..c9c3ab6cd 100644 --- a/internal/types/knowledgebase.go +++ b/internal/types/knowledgebase.go @@ -188,6 +188,7 @@ func (c *VLMConfig) Scan(value interface{}) error { } type ExtractConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` Text string `yaml:"text" json:"text"` Tags []string `yaml:"tags" json:"tags"` Nodes []*GraphNode `yaml:"nodes" json:"nodes"`