271 lines
9.6 KiB
Go
271 lines
9.6 KiB
Go
package chatpipline
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"knowlege-lsxd/internal/tracing"
|
|
"knowlege-lsxd/internal/types"
|
|
)
|
|
|
|
// PluginTracing implements tracing functionality for chat pipeline events
|
|
type PluginTracing struct{}
|
|
|
|
// NewPluginTracing creates a new tracing plugin instance
|
|
func NewPluginTracing(eventManager *EventManager) *PluginTracing {
|
|
res := &PluginTracing{}
|
|
eventManager.Register(res)
|
|
return res
|
|
}
|
|
|
|
// ActivationEvents returns the event types this plugin handles
|
|
func (p *PluginTracing) ActivationEvents() []types.EventType {
|
|
return []types.EventType{
|
|
types.CHUNK_SEARCH,
|
|
types.CHUNK_RERANK,
|
|
types.CHUNK_MERGE,
|
|
types.INTO_CHAT_MESSAGE,
|
|
types.CHAT_COMPLETION,
|
|
types.CHAT_COMPLETION_STREAM,
|
|
types.FILTER_TOP_K,
|
|
types.REWRITE_QUERY,
|
|
types.PREPROCESS_QUERY,
|
|
}
|
|
}
|
|
|
|
// OnEvent handles incoming events and routes them to the appropriate tracing handler based on event type.
|
|
// It acts as the central dispatcher for all tracing-related events in the chat pipeline.
|
|
//
|
|
// Parameters:
|
|
// - ctx: context.Context for request-scoped values, cancellation signals, and deadlines
|
|
// - eventType: the type of event being processed (e.g., CHUNK_SEARCH, CHAT_COMPLETION)
|
|
// - chatManage: contains all the chat-related data and state for the current request
|
|
// - next: callback function to continue processing in the pipeline
|
|
//
|
|
// Returns:
|
|
// - *PluginError: error if any occurred during processing, or nil if successful
|
|
func (p *PluginTracing) OnEvent(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
switch eventType {
|
|
case types.CHUNK_SEARCH:
|
|
return p.Search(ctx, eventType, chatManage, next)
|
|
case types.CHUNK_RERANK:
|
|
return p.Rerank(ctx, eventType, chatManage, next)
|
|
case types.CHUNK_MERGE:
|
|
return p.Merge(ctx, eventType, chatManage, next)
|
|
case types.INTO_CHAT_MESSAGE:
|
|
return p.IntoChatMessage(ctx, eventType, chatManage, next)
|
|
case types.CHAT_COMPLETION:
|
|
return p.ChatCompletion(ctx, eventType, chatManage, next)
|
|
case types.CHAT_COMPLETION_STREAM:
|
|
return p.ChatCompletionStream(ctx, eventType, chatManage, next)
|
|
case types.FILTER_TOP_K:
|
|
return p.FilterTopK(ctx, eventType, chatManage, next)
|
|
case types.REWRITE_QUERY:
|
|
return p.RewriteQuery(ctx, eventType, chatManage, next)
|
|
case types.PREPROCESS_QUERY:
|
|
return p.PreprocessQuery(ctx, eventType, chatManage, next)
|
|
}
|
|
return next()
|
|
}
|
|
|
|
// Search traces search operations in the chat pipeline
|
|
func (p *PluginTracing) Search(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.Search")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.String("query", chatManage.Query),
|
|
attribute.Float64("vector_threshold", chatManage.VectorThreshold),
|
|
attribute.Float64("keyword_threshold", chatManage.KeywordThreshold),
|
|
attribute.Int("match_count", chatManage.EmbeddingTopK),
|
|
)
|
|
err := next()
|
|
searchResultJson, _ := json.Marshal(chatManage.SearchResult)
|
|
span.SetAttributes(
|
|
attribute.String("hybrid_search", string(searchResultJson)),
|
|
attribute.String("processed_query", chatManage.ProcessedQuery),
|
|
)
|
|
return err
|
|
}
|
|
|
|
// Rerank traces rerank operations in the chat pipeline
|
|
func (p *PluginTracing) Rerank(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.Rerank")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.String("query", chatManage.Query),
|
|
attribute.Int("passages_count", len(chatManage.SearchResult)),
|
|
attribute.String("rerank_model_id", chatManage.RerankModelID),
|
|
attribute.Float64("rerank_filter_threshold", chatManage.RerankThreshold),
|
|
attribute.Int("rerank_filter_topk", chatManage.RerankTopK),
|
|
)
|
|
err := next()
|
|
resultJson, _ := json.Marshal(chatManage.RerankResult)
|
|
span.SetAttributes(
|
|
attribute.Int("rerank_resp_count", len(chatManage.RerankResult)),
|
|
attribute.String("rerank_resp_results", string(resultJson)),
|
|
)
|
|
return err
|
|
}
|
|
|
|
// Merge traces merge operations in the chat pipeline
|
|
func (p *PluginTracing) Merge(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.Merge")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.Int("search_results_count", len(chatManage.SearchResult)),
|
|
attribute.Int("rerank_results_count", len(chatManage.RerankResult)),
|
|
)
|
|
err := next()
|
|
mergeResultJson, _ := json.Marshal(chatManage.MergeResult)
|
|
span.SetAttributes(
|
|
attribute.Int("merge_results_count", len(chatManage.MergeResult)),
|
|
attribute.String("merge_results", string(mergeResultJson)),
|
|
)
|
|
return err
|
|
}
|
|
|
|
// IntoChatMessage traces message conversion operations
|
|
func (p *PluginTracing) IntoChatMessage(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.IntoChatMessage")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.Int("search_results_count", len(chatManage.SearchResult)),
|
|
attribute.Int("rerank_results_count", len(chatManage.RerankResult)),
|
|
attribute.Int("merge_results_count", len(chatManage.MergeResult)),
|
|
)
|
|
err := next()
|
|
span.SetAttributes(attribute.Int("generated_content_length", len(chatManage.UserContent)))
|
|
return err
|
|
}
|
|
|
|
// ChatCompletion traces chat completion operations
|
|
func (p *PluginTracing) ChatCompletion(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.ChatCompletion")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.String("model_id", chatManage.ChatModelID),
|
|
attribute.String("system_prompt", chatManage.SummaryConfig.Prompt),
|
|
attribute.String("user_prompt", chatManage.UserContent),
|
|
attribute.Int("total_references", len(chatManage.RerankResult)),
|
|
)
|
|
err := next()
|
|
span.SetAttributes(
|
|
attribute.String("chat_response", chatManage.ChatResponse.Content),
|
|
attribute.Int("chat_response_tokens", chatManage.ChatResponse.Usage.TotalTokens),
|
|
attribute.Int("chat_response_prompt_tokens", chatManage.ChatResponse.Usage.PromptTokens),
|
|
attribute.Int("chat_response_completion_tokens", chatManage.ChatResponse.Usage.CompletionTokens),
|
|
)
|
|
return err
|
|
}
|
|
|
|
// ChatCompletionStream traces streaming chat completion operations
|
|
func (p *PluginTracing) ChatCompletionStream(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
ctx, span := tracing.ContextWithSpan(ctx, "PluginTracing.ChatCompletionStream")
|
|
startTime := time.Now()
|
|
span.SetAttributes(
|
|
attribute.String("model_id", chatManage.ChatModelID),
|
|
attribute.String("system_prompt", chatManage.SummaryConfig.Prompt),
|
|
attribute.String("user_prompt", chatManage.UserContent),
|
|
attribute.Int("total_references", len(chatManage.RerankResult)),
|
|
)
|
|
|
|
responseBuilder := &strings.Builder{}
|
|
oldStream := chatManage.ResponseChan
|
|
newStream := make(chan types.StreamResponse)
|
|
chatManage.ResponseChan = newStream
|
|
|
|
go func(ctx context.Context) {
|
|
for resp := range oldStream {
|
|
if resp.ResponseType == types.ResponseTypeAnswer {
|
|
responseBuilder.WriteString(resp.Content)
|
|
}
|
|
newStream <- resp
|
|
}
|
|
elapsedMS := time.Since(startTime).Milliseconds()
|
|
span.SetAttributes(
|
|
attribute.Bool("chat_completion_success", true),
|
|
attribute.Int64("response_time_ms", elapsedMS),
|
|
attribute.String("chat_response", responseBuilder.String()),
|
|
attribute.Int("final_response_length", responseBuilder.Len()),
|
|
attribute.Float64("tokens_per_second", float64(responseBuilder.Len())/float64(elapsedMS)*1000),
|
|
)
|
|
span.End()
|
|
close(newStream)
|
|
}(ctx)
|
|
|
|
return next()
|
|
}
|
|
|
|
// FilterTopK traces filtering operations in the chat pipeline
|
|
func (p *PluginTracing) FilterTopK(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.FilterTopK")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.Int("before_filter_search_results_count", len(chatManage.SearchResult)),
|
|
attribute.Int("before_filter_rerank_results_count", len(chatManage.RerankResult)),
|
|
attribute.Int("before_filter_merge_results_count", len(chatManage.MergeResult)),
|
|
)
|
|
err := next()
|
|
span.SetAttributes(
|
|
attribute.Int("after_filter_search_results_count", len(chatManage.SearchResult)),
|
|
attribute.Int("after_filter_rerank_results_count", len(chatManage.RerankResult)),
|
|
attribute.Int("after_filter_merge_results_count", len(chatManage.MergeResult)),
|
|
)
|
|
return err
|
|
}
|
|
|
|
// RewriteQuery traces query rewriting operations
|
|
func (p *PluginTracing) RewriteQuery(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.RewriteQuery")
|
|
defer span.End()
|
|
span.SetAttributes(
|
|
attribute.String("query", chatManage.Query),
|
|
)
|
|
err := next()
|
|
span.SetAttributes(
|
|
attribute.String("rewrite_query", chatManage.RewriteQuery),
|
|
)
|
|
return err
|
|
}
|
|
|
|
// PreprocessQuery traces query preprocessing operations
|
|
func (p *PluginTracing) PreprocessQuery(ctx context.Context,
|
|
eventType types.EventType, chatManage *types.ChatManage, next func() *PluginError,
|
|
) *PluginError {
|
|
_, span := tracing.ContextWithSpan(ctx, "PluginTracing.PreprocessQuery")
|
|
defer span.End()
|
|
|
|
span.SetAttributes(
|
|
attribute.String("query", chatManage.Query),
|
|
)
|
|
|
|
err := next()
|
|
|
|
span.SetAttributes(
|
|
attribute.String("processed_query", chatManage.ProcessedQuery),
|
|
)
|
|
|
|
return err
|
|
}
|