feat: enhance retriever summary handling by adding SetRetrieverSummary method and updating trace summary in knowledge retrieval
This commit is contained in:
@@ -163,7 +163,7 @@ func (s *Service) Run(ctx context.Context, req Request) (*Summary, error) {
|
||||
QueryPreview: preview(req.UserMessage.Content, 120),
|
||||
}, strings.TrimSpace(req.UserMessage.Content)); retrieveErr == nil && retrieveResult != nil {
|
||||
summary.RetrieverCount = len(retrieveResult.Hits)
|
||||
collector.Data.Retriever.Count = len(retrieveResult.Hits)
|
||||
collector.SetRetrieverSummary(retrieveResult.TraceSummary)
|
||||
collector.Data.Retriever.Items = append(collector.Data.Retriever.Items, retrieveResult.TraceItems...)
|
||||
if strings.TrimSpace(retrieveResult.ContextText) != "" {
|
||||
messages = append(messages, schema.SystemMessage(retrieveResult.ContextText))
|
||||
|
||||
@@ -65,6 +65,22 @@ func (c *RuntimeTraceCollector) SetSkillMiddleware(enabled bool, toolName string
|
||||
c.Data.Skill.MiddlewareToolName = toolName
|
||||
}
|
||||
|
||||
func (c *RuntimeTraceCollector) SetRetrieverSummary(summary RetrieverTraceSummary) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.Data.Retriever.TopK = summary.TopK
|
||||
c.Data.Retriever.ScoreThreshold = summary.ScoreThreshold
|
||||
c.Data.Retriever.ContextMaxTokens = summary.ContextMaxTokens
|
||||
c.Data.Retriever.Count = summary.HitCount
|
||||
c.Data.Retriever.ContextCount = summary.ContextCount
|
||||
c.Data.Retriever.EmbeddingMs = summary.EmbeddingMs
|
||||
c.Data.Retriever.VectorSearchMs = summary.VectorSearchMs
|
||||
c.Data.Retriever.HydrateMs = summary.HydrateMs
|
||||
}
|
||||
|
||||
func (c *RuntimeTraceCollector) AddToolItem(item ToolTraceItem) {
|
||||
if c == nil {
|
||||
return
|
||||
|
||||
@@ -41,6 +41,17 @@ type RetrieverTraceItem struct {
|
||||
LatencyMs int64 `json:"latencyMs,omitempty"`
|
||||
}
|
||||
|
||||
type RetrieverTraceSummary struct {
|
||||
TopK int
|
||||
ScoreThreshold float64
|
||||
ContextMaxTokens int
|
||||
HitCount int
|
||||
ContextCount int
|
||||
EmbeddingMs int64
|
||||
VectorSearchMs int64
|
||||
HydrateMs int64
|
||||
}
|
||||
|
||||
type InstructionTraceSummary struct {
|
||||
SectionTitles []string
|
||||
HasProjectRule bool
|
||||
@@ -81,8 +92,15 @@ type RuntimeTraceData struct {
|
||||
CurrentUserMessagePreview string `json:"currentUserMessagePreview,omitempty"`
|
||||
} `json:"input"`
|
||||
Retriever struct {
|
||||
Count int `json:"count,omitempty"`
|
||||
Items []RetrieverTraceItem `json:"items,omitempty"`
|
||||
Count int `json:"count,omitempty"`
|
||||
TopK int `json:"topK,omitempty"`
|
||||
ScoreThreshold float64 `json:"scoreThreshold,omitempty"`
|
||||
ContextMaxTokens int `json:"contextMaxTokens,omitempty"`
|
||||
ContextCount int `json:"contextCount,omitempty"`
|
||||
EmbeddingMs int64 `json:"embeddingMs,omitempty"`
|
||||
VectorSearchMs int64 `json:"vectorSearchMs,omitempty"`
|
||||
HydrateMs int64 `json:"hydrateMs,omitempty"`
|
||||
Items []RetrieverTraceItem `json:"items,omitempty"`
|
||||
} `json:"retriever"`
|
||||
Tools struct {
|
||||
Count int `json:"count,omitempty"`
|
||||
|
||||
@@ -26,11 +26,13 @@ type KnowledgeRetrieveOptions struct {
|
||||
type KnowledgeRetrieveResult struct {
|
||||
KnowledgeBaseIDs []int64
|
||||
Query string
|
||||
Options KnowledgeRetrieveOptions
|
||||
Hits []rag.RetrieveResult
|
||||
ContextResults []rag.RetrieveResult
|
||||
ContextText string
|
||||
Trace *rag.RetrieveTrace
|
||||
TraceItems []callbacks.RetrieverTraceItem
|
||||
TraceSummary callbacks.RetrieverTraceSummary
|
||||
}
|
||||
|
||||
func NewKnowledgeRetriever(aiAgent *models.AIAgent) *KnowledgeRetriever {
|
||||
@@ -76,6 +78,12 @@ func (r *KnowledgeRetriever) RetrieveContextByOptions(ctx context.Context, opts
|
||||
ret := &KnowledgeRetrieveResult{
|
||||
KnowledgeBaseIDs: append([]int64(nil), knowledgeBaseIDs...),
|
||||
Query: query,
|
||||
Options: KnowledgeRetrieveOptions{
|
||||
ContextMaxTokens: contextMaxTokens,
|
||||
TopK: opts.TopK,
|
||||
ScoreThreshold: opts.ScoreThreshold,
|
||||
QueryPreview: queryPreview,
|
||||
},
|
||||
}
|
||||
if query == "" || len(knowledgeBaseIDs) == 0 {
|
||||
return ret, nil
|
||||
@@ -89,6 +97,7 @@ func (r *KnowledgeRetriever) RetrieveContextByOptions(ctx context.Context, opts
|
||||
ret.ContextResults = rag.Retrieve.SelectContextResults(results, contextMaxTokens)
|
||||
ret.ContextText = strings.TrimSpace(rag.Retrieve.BuildContext(ctx, results, contextMaxTokens))
|
||||
ret.TraceItems = buildRetrieverTraceItems(queryPreview, results, trace)
|
||||
ret.TraceSummary = buildRetrieverTraceSummary(ret.Options, ret.ContextResults, results, trace)
|
||||
return ret, nil
|
||||
}
|
||||
|
||||
@@ -113,3 +122,19 @@ func buildRetrieverTraceItems(queryPreview string, results []rag.RetrieveResult,
|
||||
}
|
||||
return ret
|
||||
}
|
||||
|
||||
func buildRetrieverTraceSummary(opts KnowledgeRetrieveOptions, contextResults []rag.RetrieveResult, results []rag.RetrieveResult, trace *rag.RetrieveTrace) callbacks.RetrieverTraceSummary {
|
||||
ret := callbacks.RetrieverTraceSummary{
|
||||
TopK: opts.TopK,
|
||||
ScoreThreshold: opts.ScoreThreshold,
|
||||
ContextMaxTokens: opts.ContextMaxTokens,
|
||||
HitCount: len(results),
|
||||
ContextCount: len(contextResults),
|
||||
}
|
||||
if trace != nil {
|
||||
ret.EmbeddingMs = trace.EmbeddingMs
|
||||
ret.VectorSearchMs = trace.VectorSearchMs
|
||||
ret.HydrateMs = trace.HydrateMs
|
||||
}
|
||||
return ret
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user