wlw/ai-agent
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: refactor logging and execution plan handling, introducing RunLogService for improved structure

Browse Source
This commit is contained in:
mlogclub
2026-04-13 17:27:32 +08:00
parent 03021d5f69
commit 5c944a28c2
6 changed files with 265 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/ai/rag/retrieve.go
+14 -131
View File
@@ -4,17 +4,14 @@ import (
"context"
"fmt"
"log/slog"
"sort"
"strings"
"time"
"cs-agent/internal/models"
"github.com/mlogclub/simple/sqls"
"cs-agent/internal/ai"
"cs-agent/internal/ai/rag/vectordb"
"cs-agent/internal/models"
"cs-agent/internal/pkg/enums"
"cs-agent/internal/repositories"
"github.com/mlogclub/simple/sqls"
)
type retrieve struct {
@@ -50,138 +47,24 @@ func (s *retrieve) RetrieveWithTrace(ctx context.Context, req RetrieveRequest) (
return nil, trace, nil
}
embeddingStartedAt := time.Now()
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, req.Query)
trace.EmbeddingMs = time.Since(embeddingStartedAt).Milliseconds()
searchResults, searchTrace, err := s.searchKnowledgeBaseVectors(ctx, req, retrievableKnowledgeBases)
if err != nil {
return nil, trace, fmt.Errorf("failed to generate query embedding: %w", err)
}
collectionName := knowledgeCollectionName
provider := vectordb.GetProvider()
if provider == nil {
return nil, trace, fmt.Errorf("vectordb provider not initialized")
}
searchResults := make([]vectordb.SearchResult, 0)
vectorSearchStartedAt := time.Now()
for _, knowledgeBase := range retrievableKnowledgeBases {
topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(req, &knowledgeBase)
kbResults, searchErr := provider.Search(ctx, &vectordb.SearchRequest{
CollectionName: collectionName,
Vector: embeddingResult.Vector,
TopK: topK,
ScoreThreshold: scoreThreshold,
Filter: &vectordb.SearchFilter{
KnowledgeBaseIDs: []int64{knowledgeBase.ID},
},
})
if searchErr != nil {
slog.Error("Failed to search vectors",
"knowledge_base_id", knowledgeBase.ID,
"error", searchErr)
trace.VectorSearchMs = time.Since(vectorSearchStartedAt).Milliseconds()
return nil, trace, fmt.Errorf("failed to search vectors: %w", searchErr)
if searchTrace != nil {
trace.EmbeddingMs = searchTrace.EmbeddingMs
trace.VectorSearchMs = searchTrace.VectorSearchMs
}
if len(kbResults) == 0 && scoreThreshold > 0 {
s.logEmptySearchDiagnostics(ctx, provider, collectionName, embeddingResult.Vector, topK, scoreThreshold, []int64{knowledgeBase.ID}, req)
}
searchResults = append(searchResults, kbResults...)
return nil, trace, err
}
if searchTrace != nil {
trace.EmbeddingMs = searchTrace.EmbeddingMs
trace.VectorSearchMs = searchTrace.VectorSearchMs
}
trace.VectorSearchMs = time.Since(vectorSearchStartedAt).Milliseconds()
if len(searchResults) == 0 {
return nil, trace, nil
}
sort.SliceStable(searchResults, func(i, j int) bool {
if searchResults[i].Score == searchResults[j].Score {
return searchResults[i].ID < searchResults[j].ID
}
return searchResults[i].Score > searchResults[j].Score
})
results := make([]RetrieveResult, 0, len(searchResults))
hydrateStartedAt := time.Now()
vectorIDs := make([]string, 0, len(searchResults))
for _, sr := range searchResults {
if strings.TrimSpace(sr.ID) == "" {
continue
}
vectorIDs = append(vectorIDs, sr.ID)
}
chunks := repositories.KnowledgeChunkRepository.FindByVectorIDs(sqls.DB(), vectorIDs)
chunkByVectorID := make(map[string]*models.KnowledgeChunk, len(chunks))
documentIDs := make([]int64, 0)
faqIDs := make([]int64, 0)
documentSeen := make(map[int64]struct{})
faqSeen := make(map[int64]struct{})
for i := range chunks {
chunk := &chunks[i]
chunkByVectorID[chunk.VectorID] = chunk
if chunk.DocumentID > 0 {
if _, ok := documentSeen[chunk.DocumentID]; !ok {
documentSeen[chunk.DocumentID] = struct{}{}
documentIDs = append(documentIDs, chunk.DocumentID)
}
}
if chunk.FaqID > 0 {
if _, ok := faqSeen[chunk.FaqID]; !ok {
faqSeen[chunk.FaqID] = struct{}{}
faqIDs = append(faqIDs, chunk.FaqID)
}
}
}
documents := repositories.KnowledgeDocumentRepository.FindByIDs(sqls.DB(), documentIDs)
documentByID := make(map[int64]*models.KnowledgeDocument, len(documents))
for i := range documents {
document := &documents[i]
documentByID[document.ID] = document
}
faqs := repositories.KnowledgeFAQRepository.FindByIDs(sqls.DB(), faqIDs)
faqByID := make(map[int64]*models.KnowledgeFAQ, len(faqs))
for i := range faqs {
faq := &faqs[i]
faqByID[faq.ID] = faq
}
for _, sr := range searchResults {
chunk := chunkByVectorID[sr.ID]
if chunk == nil || chunk.Status != enums.StatusOk {
continue
}
documentTitle := ""
faqQuestion := ""
if chunk.DocumentID > 0 {
document := documentByID[chunk.DocumentID]
if document == nil || document.Status != enums.StatusOk {
continue
}
documentTitle = document.Title
}
if chunk.FaqID > 0 {
faq := faqByID[chunk.FaqID]
if faq == nil || faq.Status != enums.StatusOk {
continue
}
faqQuestion = faq.Question
}
results = append(results, RetrieveResult{
KnowledgeBaseID: chunk.KnowledgeBaseID,
ChunkID: chunk.ID,
DocumentID: chunk.DocumentID,
DocumentTitle: documentTitle,
FaqID: chunk.FaqID,
FaqQuestion: faqQuestion,
ChunkNo: chunk.ChunkNo,
Title: chunk.Title,
SectionPath: chunk.SectionPath,
Content: chunk.Content,
Score: sr.Score,
ChunkType: extractChunkType(sr.Payload),
})
}
trace.HydrateMs = time.Since(hydrateStartedAt).Milliseconds()
results, hydrateMs := s.hydrateRetrieveResults(searchResults)
trace.HydrateMs = hydrateMs
return results, trace, nil
}
internal/ai/rag/retrieve_search.go
+163
View File
@@ -0,0 +1,163 @@
package rag
import (
"context"
"fmt"
"log/slog"
"sort"
"strings"
"time"
"cs-agent/internal/ai"
"cs-agent/internal/ai/rag/vectordb"
"cs-agent/internal/models"
"cs-agent/internal/pkg/enums"
"cs-agent/internal/repositories"
"github.com/mlogclub/simple/sqls"
)
func (s *retrieve) searchKnowledgeBaseVectors(ctx context.Context, req RetrieveRequest, knowledgeBases []models.KnowledgeBase) ([]vectordb.SearchResult, *RetrieveTrace, error) {
trace := &RetrieveTrace{}
embeddingStartedAt := time.Now()
embeddingResult, err := ai.Embedding.GenerateEmbedding(ctx, req.Query)
trace.EmbeddingMs = time.Since(embeddingStartedAt).Milliseconds()
if err != nil {
return nil, trace, fmt.Errorf("failed to generate query embedding: %w", err)
}
collectionName := knowledgeCollectionName
provider := vectordb.GetProvider()
if provider == nil {
return nil, trace, fmt.Errorf("vectordb provider not initialized")
}
searchResults := make([]vectordb.SearchResult, 0)
vectorSearchStartedAt := time.Now()
for _, knowledgeBase := range knowledgeBases {
topK, scoreThreshold := resolveKnowledgeBaseSearchOptions(req, &knowledgeBase)
kbResults, searchErr := provider.Search(ctx, &vectordb.SearchRequest{
CollectionName: collectionName,
Vector: embeddingResult.Vector,
TopK: topK,
ScoreThreshold: scoreThreshold,
Filter: &vectordb.SearchFilter{
KnowledgeBaseIDs: []int64{knowledgeBase.ID},
},
})
if searchErr != nil {
slog.Error("Failed to search vectors",
"knowledge_base_id", knowledgeBase.ID,
"error", searchErr)
trace.VectorSearchMs = time.Since(vectorSearchStartedAt).Milliseconds()
return nil, trace, fmt.Errorf("failed to search vectors: %w", searchErr)
}
if len(kbResults) == 0 && scoreThreshold > 0 {
s.logEmptySearchDiagnostics(ctx, provider, collectionName, embeddingResult.Vector, topK, scoreThreshold, []int64{knowledgeBase.ID}, req)
}
searchResults = append(searchResults, kbResults...)
}
trace.VectorSearchMs = time.Since(vectorSearchStartedAt).Milliseconds()
if len(searchResults) > 0 {
sort.SliceStable(searchResults, func(i, j int) bool {
if searchResults[i].Score == searchResults[j].Score {
return searchResults[i].ID < searchResults[j].ID
}
return searchResults[i].Score > searchResults[j].Score
})
}
return searchResults, trace, nil
}
func (s *retrieve) hydrateRetrieveResults(searchResults []vectordb.SearchResult) ([]RetrieveResult, int64) {
if len(searchResults) == 0 {
return nil, 0
}
hydrateStartedAt := time.Now()
results := make([]RetrieveResult, 0, len(searchResults))
vectorIDs := make([]string, 0, len(searchResults))
for _, sr := range searchResults {
if strings.TrimSpace(sr.ID) == "" {
continue
}
vectorIDs = append(vectorIDs, sr.ID)
}
chunks := repositories.KnowledgeChunkRepository.FindByVectorIDs(sqls.DB(), vectorIDs)
chunkByVectorID := make(map[string]*models.KnowledgeChunk, len(chunks))
documentIDs := make([]int64, 0)
faqIDs := make([]int64, 0)
documentSeen := make(map[int64]struct{})
faqSeen := make(map[int64]struct{})
for i := range chunks {
chunk := &chunks[i]
chunkByVectorID[chunk.VectorID] = chunk
if chunk.DocumentID > 0 {
if _, ok := documentSeen[chunk.DocumentID]; !ok {
documentSeen[chunk.DocumentID] = struct{}{}
documentIDs = append(documentIDs, chunk.DocumentID)
}
}
if chunk.FaqID > 0 {
if _, ok := faqSeen[chunk.FaqID]; !ok {
faqSeen[chunk.FaqID] = struct{}{}
faqIDs = append(faqIDs, chunk.FaqID)
}
}
}
documents := repositories.KnowledgeDocumentRepository.FindByIDs(sqls.DB(), documentIDs)
documentByID := make(map[int64]*models.KnowledgeDocument, len(documents))
for i := range documents {
document := &documents[i]
documentByID[document.ID] = document
}
faqs := repositories.KnowledgeFAQRepository.FindByIDs(sqls.DB(), faqIDs)
faqByID := make(map[int64]*models.KnowledgeFAQ, len(faqs))
for i := range faqs {
faq := &faqs[i]
faqByID[faq.ID] = faq
}
for _, sr := range searchResults {
chunk := chunkByVectorID[sr.ID]
if chunk == nil || chunk.Status != enums.StatusOk {
continue
}
documentTitle := ""
faqQuestion := ""
if chunk.DocumentID > 0 {
document := documentByID[chunk.DocumentID]
if document == nil || document.Status != enums.StatusOk {
continue
}
documentTitle = document.Title
}
if chunk.FaqID > 0 {
faq := faqByID[chunk.FaqID]
if faq == nil || faq.Status != enums.StatusOk {
continue
}
faqQuestion = faq.Question
}
results = append(results, RetrieveResult{
KnowledgeBaseID: chunk.KnowledgeBaseID,
ChunkID: chunk.ID,
DocumentID: chunk.DocumentID,
DocumentTitle: documentTitle,
FaqID: chunk.FaqID,
FaqQuestion: faqQuestion,
ChunkNo: chunk.ChunkNo,
Title: chunk.Title,
SectionPath: chunk.SectionPath,
Content: chunk.Content,
Score: sr.Score,
ChunkType: extractChunkType(sr.Payload),
})
}
return results, time.Since(hydrateStartedAt).Milliseconds()
}
internal/ai/skills/log.go
+5 -48
View File
@@ -1,56 +1,13 @@
package skills
import (
"encoding/json"
"time"
"cs-agent/internal/models"
)
import "cs-agent/internal/models"
// BuildRunLog 根据执行计划与运行结果构建 Skill 运行日志。
func BuildRunLog(ctx RuntimeContext, plan *ExecutionPlan, trace *ExecutionTrace, err error) *models.SkillRunLog {
log := &models.SkillRunLog{
ConversationID: ctx.ConversationID,
AIAgentID: ctx.AIAgentID,
ManualSkillCode: ctx.ManualSkillCode,
IntentCode: ctx.IntentCode,
UserMessage: ctx.UserMessage,
TraceData: buildTraceData(trace),
CreatedAt: time.Now(),
}
if plan != nil {
if plan.AIConfig != nil {
log.AIConfigID = plan.AIConfig.ID
log.UsedModel = plan.AIConfig.ModelName
log.UsedProvider = plan.AIConfig.Provider
}
if plan.Skill != nil {
log.SkillDefinitionID = plan.Skill.ID
log.SkillCode = plan.Skill.Code
log.Matched = true
log.FinalSelected = true
log.MatchReason = plan.MatchReason
}
}
if err != nil {
log.ErrorMessage = err.Error()
} else if !log.Matched {
if plan != nil && plan.MatchReason != "" {
log.MatchReason = plan.MatchReason
} else {
log.MatchReason = "not_matched"
}
}
return log
return RuntimeService.runlog.Build(ctx, plan, trace, err)
}
func buildTraceData(trace *ExecutionTrace) string {
if trace == nil {
return ""
}
data, err := json.Marshal(trace)
if err != nil {
return ""
}
return string(data)
// WriteRunLog 写入 Skill 路由日志。
func WriteRunLog(log *models.SkillRunLog) error {
return RuntimeService.runlog.Write(log)
}
internal/ai/skills/runlog_service.go
+73
View File
@@ -0,0 +1,73 @@
package skills
import (
"encoding/json"
"time"
"cs-agent/internal/models"
"cs-agent/internal/repositories"
"github.com/mlogclub/simple/sqls"
)
func newRunLogService() *RunLogService {
return &RunLogService{}
}
type RunLogService struct{}
// Build 根据执行计划与运行结果构建 Skill 运行日志。
func (s *RunLogService) Build(ctx RuntimeContext, plan *ExecutionPlan, trace *ExecutionTrace, err error) *models.SkillRunLog {
log := &models.SkillRunLog{
ConversationID: ctx.ConversationID,
AIAgentID: ctx.AIAgentID,
ManualSkillCode: ctx.ManualSkillCode,
IntentCode: ctx.IntentCode,
UserMessage: ctx.UserMessage,
TraceData: s.buildTraceData(trace),
CreatedAt: time.Now(),
}
if plan != nil {
if plan.AIConfig != nil {
log.AIConfigID = plan.AIConfig.ID
log.UsedModel = plan.AIConfig.ModelName
log.UsedProvider = plan.AIConfig.Provider
}
if plan.Skill != nil {
log.SkillDefinitionID = plan.Skill.ID
log.SkillCode = plan.Skill.Code
log.Matched = true
log.FinalSelected = true
log.MatchReason = plan.MatchReason
}
}
if err != nil {
log.ErrorMessage = err.Error()
} else if !log.Matched {
if plan != nil && plan.MatchReason != "" {
log.MatchReason = plan.MatchReason
} else {
log.MatchReason = "not_matched"
}
}
return log
}
// Write 写入 Skill 路由日志。
func (s *RunLogService) Write(log *models.SkillRunLog) error {
if log == nil {
return nil
}
return repositories.SkillRunLogRepository.Create(sqls.DB(), log)
}
func (s *RunLogService) buildTraceData(trace *ExecutionTrace) string {
if trace == nil {
return ""
}
data, err := json.Marshal(trace)
if err != nil {
return ""
}
return string(data)
}
internal/ai/skills/runtime.go
-7
View File
@@ -2,8 +2,6 @@ package skills
import (
"context"
"cs-agent/internal/models"
)
// BuildExecutionPlan 构建当前请求的 Skill 执行计划。
@@ -11,11 +9,6 @@ func BuildExecutionPlan(execCtx context.Context, ctx RuntimeContext) (*Execution
return RuntimeService.BuildExecutionPlan(execCtx, ctx)
}
// WriteRunLog 写入 Skill 路由日志。
func WriteRunLog(log *models.SkillRunLog) error {
return RuntimeService.WriteRunLog(log)
}
// Select 执行一次 Skill 路由并记录路由日志。
func Select(ctx context.Context, runtimeCtx RuntimeContext) (*ExecutionResult, error) {
return RuntimeService.Select(ctx, runtimeCtx)
internal/ai/skills/service.go
+10 -9
View File
@@ -14,10 +14,14 @@ import (
var RuntimeService = newService()
func newService() *Service {
return &Service{}
return &Service{
runlog: newRunLogService(),
}
}
type Service struct{}
type Service struct {
runlog *RunLogService
}
// BuildExecutionPlan 构建当前请求的 Skill 执行计划。
func (s *Service) BuildExecutionPlan(execCtx context.Context, ctx RuntimeContext) (*ExecutionPlan, error) {
@@ -50,10 +54,7 @@ func (s *Service) BuildExecutionPlan(execCtx context.Context, ctx RuntimeContext
// WriteRunLog 写入 Skill 路由日志。
func (s *Service) WriteRunLog(log *models.SkillRunLog) error {
if log == nil {
return nil
}
return repositories.SkillRunLogRepository.Create(sqls.DB(), log)
return s.runlog.Write(log)
}
// Select 执行一次 Skill 路由并记录路由日志。
@@ -61,7 +62,7 @@ func (s *Service) Select(ctx context.Context, runtimeCtx RuntimeContext) (*Execu
plan, err := s.BuildExecutionPlan(ctx, runtimeCtx)
if err != nil {
trace := &ExecutionTrace{Status: "route_error"}
log := BuildRunLog(runtimeCtx, nil, trace, err)
log := s.runlog.Build(runtimeCtx, nil, trace, err)
_ = s.WriteRunLog(log)
return nil, err
}
@@ -72,7 +73,7 @@ func (s *Service) Select(ctx context.Context, runtimeCtx RuntimeContext) (*Execu
trace.MatchReason = strings.TrimSpace(plan.MatchReason)
trace.Route = plan.RouteTrace
}
log := BuildRunLog(runtimeCtx, plan, trace, nil)
log := s.runlog.Build(runtimeCtx, plan, trace, nil)
_ = s.WriteRunLog(log)
return &ExecutionResult{
Plan: plan,
@@ -82,7 +83,7 @@ func (s *Service) Select(ctx context.Context, runtimeCtx RuntimeContext) (*Execu
}
trace.MatchReason = strings.TrimSpace(plan.MatchReason)
trace.Route = plan.RouteTrace
log := BuildRunLog(runtimeCtx, plan, trace, err)
log := s.runlog.Build(runtimeCtx, plan, trace, err)
if writeErr := s.WriteRunLog(log); writeErr != nil && err == nil {
err = writeErr
}