diff --git a/internal/application/service/chat_pipeline/rerank.go b/internal/application/service/chat_pipeline/rerank.go index 20ff2348..d33fe831 100644 --- a/internal/application/service/chat_pipeline/rerank.go +++ b/internal/application/service/chat_pipeline/rerank.go @@ -10,6 +10,7 @@ import ( "github.com/Tencent/WeKnora/internal/models/rerank" "github.com/Tencent/WeKnora/internal/searchutil" + "github.com/Tencent/WeKnora/internal/tracing/langfuse" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) @@ -83,9 +84,7 @@ func (p *PluginRerank) OnEvent(ctx context.Context, }) continue } - // 合并Content和ImageInfo的文本内容 passage := getEnrichedPassage(ctx, result) - // Skip passages that become empty after cleaning if strings.TrimSpace(passage) == "" { pipelineInfo(ctx, "Rerank", "empty_passage_skip", map[string]interface{}{ "chunk_id": result.ID, @@ -96,6 +95,31 @@ func (p *PluginRerank) OnEvent(ctx context.Context, candidatesToRerank = append(candidatesToRerank, result) } + passagesPreview := langfuse.SummarizePassagePreviews(candidatesToRerank, passages, 25) + rerankCtx, rerankSpan := langfuse.GetManager().StartSpan(ctx, langfuse.SpanOptions{ + Name: "rerank", + Input: map[string]interface{}{ + "query": chatManage.RewriteQuery, + "candidate_count": len(candidatesToRerank), + "direct_load_count": len(directLoadResults), + "rerank_model_id": chatManage.RerankModelID, + "threshold": chatManage.RerankThreshold, + "rerank_top_k": chatManage.RerankTopK, + "faq_priority": chatManage.FAQPriorityEnabled, + "faq_score_boost": chatManage.FAQScoreBoost, + "passages_preview": passagesPreview, + }, + Metadata: map[string]interface{}{ + "session_id": chatManage.SessionID, + }, + }) + ctx = rerankCtx + spanOutput := map[string]interface{}{} + var spanErr error + defer func() { + rerankSpan.Finish(spanOutput, nil, spanErr) + }() + pipelineInfo(ctx, "Rerank", "build_passages", map[string]interface{}{ "total_cnt": len(chatManage.SearchResult), "candidate_cnt": len(candidatesToRerank), @@ -103,6 +127,8 @@ func (p *PluginRerank) OnEvent(ctx context.Context, }) var rerankResp []rerank.RankResult + var rawRerankResp []rerank.RankResult + thresholdDegraded := false // Only call rerank model if there are candidates if len(candidatesToRerank) > 0 { @@ -119,11 +145,18 @@ func (p *PluginRerank) OnEvent(ctx context.Context, "candidate_cnt": len(candidatesToRerank), }) chatManage.SearchResult = append(directLoadResults, candidatesToRerank...) + spanOutput = map[string]interface{}{ + "stage": "api_error_fallback", + "candidate_count": len(candidatesToRerank), + "error": rerankErr.Error(), + } return next() } + rawRerankResp = append([]rerank.RankResult(nil), rerankResp...) // If no results and threshold is high enough, try with lower threshold if len(rerankResp) == 0 && originalThreshold > 0.3 { + thresholdDegraded = true degradedThreshold := originalThreshold * 0.7 if degradedThreshold < 0.3 { degradedThreshold = 0.3 @@ -144,8 +177,15 @@ func (p *PluginRerank) OnEvent(ctx context.Context, "candidate_cnt": len(candidatesToRerank), }) chatManage.SearchResult = append(directLoadResults, candidatesToRerank...) + spanOutput = map[string]interface{}{ + "stage": "api_error_fallback", + "candidate_count": len(candidatesToRerank), + "threshold_degraded": thresholdDegraded, + "error": rerankErr.Error(), + } return next() } + rawRerankResp = append([]rerank.RankResult(nil), rerankResp...) } } @@ -169,6 +209,7 @@ func (p *PluginRerank) OnEvent(ctx context.Context, base := sr.Score sr.Metadata["base_score"] = fmt.Sprintf("%.4f", base) modelScore := rr.RelevanceScore + sr.Metadata["model_score"] = fmt.Sprintf("%.4f", modelScore) sr.Score = compositeScore(sr, modelScore, base) // Apply FAQ score boost if enabled @@ -193,8 +234,9 @@ func (p *PluginRerank) OnEvent(ctx context.Context, for _, sr := range directLoadResults { base := sr.Score sr.Metadata["base_score"] = fmt.Sprintf("%.4f", base) - // Assign high model score for direct load items modelScore := 1.0 + sr.Metadata["model_score"] = fmt.Sprintf("%.4f", modelScore) + // Assign high model score for direct load items sr.Score = compositeScore(sr, modelScore, base) reranked = append(reranked, sr) } @@ -216,15 +258,84 @@ func (p *PluginRerank) OnEvent(ctx context.Context, pipelineWarn(ctx, "Rerank", "output", map[string]interface{}{ "filtered_cnt": 0, }) + spanOutput = buildRerankSpanOutput( + candidatesToRerank, + passages, + directLoadResults, + rawRerankResp, + reranked, + nil, + chatManage, + thresholdDegraded, + ) return ErrSearchNothing } + spanOutput = buildRerankSpanOutput( + candidatesToRerank, + passages, + directLoadResults, + rawRerankResp, + reranked, + chatManage.RerankResult, + chatManage, + thresholdDegraded, + ) pipelineInfo(ctx, "Rerank", "output", map[string]interface{}{ "filtered_cnt": len(chatManage.RerankResult), }) return next() } +func buildRerankSpanOutput( + candidates []*types.SearchResult, + passages []string, + directLoad []*types.SearchResult, + modelScores []rerank.RankResult, + composite []*types.SearchResult, + final []*types.SearchResult, + chatManage *types.ChatManage, + thresholdDegraded bool, +) map[string]interface{} { + modelRows := make([]map[string]interface{}, 0, len(modelScores)) + for i, rr := range modelScores { + row := map[string]interface{}{ + "rank": i + 1, + "index": rr.Index, + "model_score": rr.RelevanceScore, + } + if rr.Index >= 0 && rr.Index < len(candidates) { + row["chunk_id"] = candidates[rr.Index].ID + row["knowledge_id"] = candidates[rr.Index].KnowledgeID + row["knowledge_title"] = candidates[rr.Index].KnowledgeTitle + row["match_type"] = candidates[rr.Index].MatchType + row["retrieval_score"] = candidates[rr.Index].Score + if rr.Index < len(passages) { + row["preview"] = langfuse.TruncateRunes(passages[rr.Index], 160) + } + } + modelRows = append(modelRows, row) + } + + out := map[string]interface{}{ + "candidate_count": len(candidates), + "direct_load_count": len(directLoad), + "model_result_count": len(modelScores), + "composite_count": len(composite), + "final_count": len(final), + "threshold": chatManage.RerankThreshold, + "rerank_top_k": chatManage.RerankTopK, + "threshold_degraded": thresholdDegraded, + "model_scores": langfuse.SummarizeRankScores(modelRows, 50), + "composite_results": langfuse.SummarizeSearchResults(composite, 25), + "final_results": langfuse.SummarizeSearchResults(final, 25), + } + if len(modelScores) > 50 { + out["model_scores_truncated"] = len(modelScores) - 50 + } + return out +} + // rerank performs the actual reranking operation with given query and passages func (p *PluginRerank) rerank(ctx context.Context, chatManage *types.ChatManage, rerankModel rerank.Reranker, query string, passages []string, diff --git a/internal/application/service/knowledgebase_search.go b/internal/application/service/knowledgebase_search.go index 605b442f..a2a42d5e 100644 --- a/internal/application/service/knowledgebase_search.go +++ b/internal/application/service/knowledgebase_search.go @@ -191,15 +191,26 @@ func (s *knowledgeBaseService) HybridSearch(ctx context.Context, retrieveCtx, retrieveSpan := langfuse.GetManager().StartSpan(ctx, langfuse.SpanOptions{ Name: "retrieve", Input: map[string]interface{}{ - "kb_ids": searchKBIDs, - "group_count": len(groups), - "match_count": matchCount, + "query_text": params.QueryText, + "kb_ids": searchKBIDs, + "knowledge_ids": params.KnowledgeIDs, + "tag_ids": params.TagIDs, + "match_count": matchCount, + "vector_threshold": params.VectorThreshold, + "keyword_threshold": params.KeywordThreshold, + "disable_vector_match": params.DisableVectorMatch, + "disable_keywords_match": params.DisableKeywordsMatch, + "group_count": len(groups), + }, + Metadata: map[string]interface{}{ + "primary_kb_id": kb.ID, + "primary_kb_type": string(kb.Type), + "embedding_model_id": kb.EmbeddingModelID, + "has_query_embedding": len(params.QueryEmbedding) > 0, }, }) retrieveResults, err := s.retrieveFromStores(retrieveCtx, groups, retriever.EngineAwareNormalizer{}) - retrieveSpan.Finish(map[string]interface{}{ - "result_count": totalHits(retrieveResults), - }, nil, err) + retrieveSpan.Finish(langfuse.SummarizeRetrieveOutput(retrieveResults), nil, err) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ "knowledge_base_ids": searchKBIDs, @@ -278,7 +289,7 @@ func allBaseParamsEmpty(groups []*storeGroup) bool { } // totalHits counts the IndexWithScore entries across a slice of retrieve -// results. Used only for langfuse span metadata. +// results. func totalHits(rrs []*types.RetrieveResult) int { n := 0 for _, r := range rrs { diff --git a/internal/models/rerank/langfuse_wrapper.go b/internal/models/rerank/langfuse_wrapper.go index ba8fa78c..f62e7856 100644 --- a/internal/models/rerank/langfuse_wrapper.go +++ b/internal/models/rerank/langfuse_wrapper.go @@ -6,6 +6,9 @@ import ( "github.com/Tencent/WeKnora/internal/tracing/langfuse" ) +const langfuseRerankPreviewDocs = 8 +const langfuseRerankMaxScores = 50 + // langfuseReranker wraps a Reranker and reports each rerank call as a // Langfuse generation observation. Rerankers don't return token usage, but // the call still incurs cost (billed per 1K documents by most vendors); we @@ -24,32 +27,75 @@ func (l *langfuseReranker) Rerank(ctx context.Context, query string, documents [ return l.inner.Rerank(ctx, query, documents) } + totalChars := len([]rune(query)) + for _, doc := range documents { + totalChars += len([]rune(doc)) + } + genCtx, gen := mgr.StartGeneration(ctx, langfuse.GenerationOptions{ Name: "rerank", Model: l.inner.GetModelName(), Input: map[string]interface{}{ - "query": query, - "document_count": len(documents), - // Only send short previews — reranker inputs can be hundreds of - // passages totalling hundreds of KB, which would bloat traces. - "documents_preview": previewDocs(documents, 5), + "query": query, + "document_count": len(documents), + "documents_preview": previewDocs(documents, langfuseRerankPreviewDocs), }, Metadata: map[string]interface{}{ - "model_id": l.inner.GetModelID(), - "num_queries": 1, + "model_id": l.inner.GetModelID(), + "num_queries": 1, + "total_chars": totalChars, + "avg_doc_chars": avgDocChars(documents), }, }) results, err := l.inner.Rerank(genCtx, query, documents) output := map[string]interface{}{ - "results": summarizeResults(results, 10), + "results": summarizeResults(results, documents, langfuseRerankMaxScores), "total_count": len(results), + "score_stats": scoreStats(results), + } + if len(results) > langfuseRerankMaxScores { + output["truncated"] = len(results) - langfuseRerankMaxScores } gen.Finish(output, approxRerankUsage(query, documents), err) return results, err } +func avgDocChars(documents []string) int { + if len(documents) == 0 { + return 0 + } + total := 0 + for _, doc := range documents { + total += len([]rune(doc)) + } + return total / len(documents) +} + +func scoreStats(results []RankResult) map[string]interface{} { + if len(results) == 0 { + return nil + } + minScore := results[0].RelevanceScore + maxScore := results[0].RelevanceScore + sum := 0.0 + for _, r := range results { + if r.RelevanceScore < minScore { + minScore = r.RelevanceScore + } + if r.RelevanceScore > maxScore { + maxScore = r.RelevanceScore + } + sum += r.RelevanceScore + } + return map[string]interface{}{ + "min": minScore, + "max": maxScore, + "avg": sum / float64(len(results)), + } +} + func approxRerankUsage(query string, documents []string) *langfuse.TokenUsage { total := len([]rune(query))/4 + 1 for _, d := range documents { @@ -80,16 +126,22 @@ func previewDocs(docs []string, n int) []map[string]interface{} { return out } -func summarizeResults(results []RankResult, n int) []map[string]interface{} { +func summarizeResults(results []RankResult, documents []string, n int) []map[string]interface{} { if len(results) < n { n = len(results) } out := make([]map[string]interface{}, 0, n) for i := 0; i < n; i++ { - out = append(out, map[string]interface{}{ - "index": results[i].Index, - "score": results[i].RelevanceScore, - }) + row := map[string]interface{}{ + "rank": i + 1, + "index": results[i].Index, + "model_score": results[i].RelevanceScore, + } + idx := results[i].Index + if idx >= 0 && idx < len(documents) { + row["preview"] = truncateRunes(documents[idx], 160) + } + out = append(out, row) } return out } diff --git a/internal/tracing/langfuse/retrieval_obs.go b/internal/tracing/langfuse/retrieval_obs.go new file mode 100644 index 00000000..cfe38b33 --- /dev/null +++ b/internal/tracing/langfuse/retrieval_obs.go @@ -0,0 +1,204 @@ +package langfuse + +import ( + "fmt" + "sort" + + "github.com/Tencent/WeKnora/internal/types" +) + +const defaultHitPreviewLimit = 25 + +// TruncateRunes shortens s to at most maxRunes runes for trace payloads. +func TruncateRunes(s string, maxRunes int) string { + if maxRunes <= 0 { + return "" + } + r := []rune(s) + if len(r) <= maxRunes { + return s + } + return string(r[:maxRunes]) + "..." +} + +// SummarizeRetrieveOutput builds Langfuse output for the retrieve span. +func SummarizeRetrieveOutput(results []*types.RetrieveResult) map[string]interface{} { + out := map[string]interface{}{ + "total_hits": 0, + "vector_hits": 0, + "keyword_hits": 0, + "group_count": len(results), + "by_retriever": []map[string]interface{}{}, + "top_hits": []map[string]interface{}{}, + } + if len(results) == 0 { + return out + } + + byRetriever := make([]map[string]interface{}, 0, len(results)) + var all []*types.IndexWithScore + for _, rr := range results { + if rr == nil { + continue + } + count := len(rr.Results) + out["total_hits"] = out["total_hits"].(int) + count + if rr.RetrieverType == types.VectorRetrieverType { + out["vector_hits"] = out["vector_hits"].(int) + count + } else { + out["keyword_hits"] = out["keyword_hits"].(int) + count + } + byRetriever = append(byRetriever, map[string]interface{}{ + "engine": string(rr.RetrieverEngineType), + "retriever": string(rr.RetrieverType), + "count": count, + }) + all = append(all, rr.Results...) + } + out["by_retriever"] = byRetriever + out["top_hits"] = summarizeIndexHits(all, defaultHitPreviewLimit) + return out +} + +// SummarizeSearchResults builds a compact ranked preview for rerank spans. +func SummarizeSearchResults(results []*types.SearchResult, limit int) map[string]interface{} { + if limit <= 0 { + limit = defaultHitPreviewLimit + } + out := map[string]interface{}{ + "count": len(results), + "top_hits": []map[string]interface{}{}, + } + if len(results) == 0 { + return out + } + + sorted := append([]*types.SearchResult(nil), results...) + sort.Slice(sorted, func(i, j int) bool { + if sorted[i].Score != sorted[j].Score { + return sorted[i].Score > sorted[j].Score + } + return sorted[i].ID < sorted[j].ID + }) + + hits := make([]map[string]interface{}, 0, minInt(limit, len(sorted))) + for i, sr := range sorted { + if i >= limit { + break + } + item := map[string]interface{}{ + "rank": i + 1, + "chunk_id": sr.ID, + "knowledge_id": sr.KnowledgeID, + "knowledge_title": sr.KnowledgeTitle, + "composite_score": fmt.Sprintf("%.4f", sr.Score), + "match_type": sr.MatchType, + "chunk_type": sr.ChunkType, + "preview": TruncateRunes(sr.Content, 160), + } + if sr.Metadata != nil { + if base, ok := sr.Metadata["base_score"]; ok { + item["retrieval_score"] = base + } + if model, ok := sr.Metadata["model_score"]; ok { + item["model_score"] = model + } + if boosted, ok := sr.Metadata["faq_boosted"]; ok { + item["faq_boosted"] = boosted + } + if orig, ok := sr.Metadata["faq_original_score"]; ok { + item["faq_original_score"] = orig + } + } + hits = append(hits, item) + } + out["top_hits"] = hits + if len(sorted) > limit { + out["truncated"] = len(sorted) - limit + } + return out +} + +// SummarizeRankScores builds rerank model score rows for Langfuse output. +func SummarizeRankScores( + results []map[string]interface{}, + limit int, +) []map[string]interface{} { + if limit <= 0 { + limit = defaultHitPreviewLimit + } + if len(results) <= limit { + return results + } + out := make([]map[string]interface{}, limit) + copy(out, results[:limit]) + return out +} + +func summarizeIndexHits(hits []*types.IndexWithScore, limit int) []map[string]interface{} { + if len(hits) == 0 { + return nil + } + sorted := append([]*types.IndexWithScore(nil), hits...) + sort.Slice(sorted, func(i, j int) bool { + if sorted[i].Score != sorted[j].Score { + return sorted[i].Score > sorted[j].Score + } + return sorted[i].ChunkID < sorted[j].ChunkID + }) + + n := minInt(limit, len(sorted)) + out := make([]map[string]interface{}, 0, n) + for i := 0; i < n; i++ { + hit := sorted[i] + out = append(out, map[string]interface{}{ + "rank": i + 1, + "chunk_id": hit.ChunkID, + "knowledge_id": hit.KnowledgeID, + "knowledge_base_id": hit.KnowledgeBaseID, + "score": fmt.Sprintf("%.4f", hit.Score), + "match_type": hit.MatchType, + "preview": TruncateRunes(hit.Content, 160), + }) + } + return out +} + +// SummarizePassagePreviews builds rerank passage previews aligned with candidates. +func SummarizePassagePreviews( + candidates []*types.SearchResult, + passages []string, + limit int, +) []map[string]interface{} { + if limit <= 0 { + limit = defaultHitPreviewLimit + } + n := len(candidates) + if len(passages) < n { + n = len(passages) + } + if limit < n { + n = limit + } + out := make([]map[string]interface{}, 0, n) + for i := 0; i < n; i++ { + sr := candidates[i] + out = append(out, map[string]interface{}{ + "index": i, + "chunk_id": sr.ID, + "knowledge_id": sr.KnowledgeID, + "knowledge_title": sr.KnowledgeTitle, + "retrieval_score": fmt.Sprintf("%.4f", sr.Score), + "match_type": sr.MatchType, + "preview": TruncateRunes(passages[i], 160), + }) + } + return out +} + +func minInt(a, b int) int { + if a < b { + return a + } + return b +} diff --git a/internal/tracing/langfuse/retrieval_obs_test.go b/internal/tracing/langfuse/retrieval_obs_test.go new file mode 100644 index 00000000..a9e2913f --- /dev/null +++ b/internal/tracing/langfuse/retrieval_obs_test.go @@ -0,0 +1,49 @@ +package langfuse + +import ( + "testing" + + "github.com/Tencent/WeKnora/internal/types" +) + +func TestSummarizeRetrieveOutput_countsAndPreviews(t *testing.T) { + out := SummarizeRetrieveOutput([]*types.RetrieveResult{ + { + RetrieverEngineType: types.PostgresRetrieverEngineType, + RetrieverType: types.VectorRetrieverType, + Results: []*types.IndexWithScore{ + {ChunkID: "c1", KnowledgeID: "k1", Score: 0.9, Content: "alpha"}, + {ChunkID: "c2", KnowledgeID: "k1", Score: 0.5, Content: "beta"}, + }, + }, + { + RetrieverEngineType: types.PostgresRetrieverEngineType, + RetrieverType: types.KeywordsRetrieverType, + Results: []*types.IndexWithScore{ + {ChunkID: "c3", KnowledgeID: "k2", Score: 0.7, Content: "gamma"}, + }, + }, + }) + + if out["total_hits"] != 3 { + t.Fatalf("total_hits = %v, want 3", out["total_hits"]) + } + if out["vector_hits"] != 2 || out["keyword_hits"] != 1 { + t.Fatalf("vector/keyword hits = %v/%v", out["vector_hits"], out["keyword_hits"]) + } + hits := out["top_hits"].([]map[string]interface{}) + if len(hits) != 3 || hits[0]["chunk_id"] != "c1" { + t.Fatalf("unexpected top_hits: %#v", hits) + } +} + +func TestSummarizeSearchResults_sortsByScore(t *testing.T) { + out := SummarizeSearchResults([]*types.SearchResult{ + {ID: "low", Score: 0.2, Content: "low"}, + {ID: "high", Score: 0.9, Content: "high"}, + }, 10) + hits := out["top_hits"].([]map[string]interface{}) + if len(hits) != 2 || hits[0]["chunk_id"] != "high" { + t.Fatalf("unexpected hits: %#v", hits) + } +}