wizardchen
beca2b89a3
Previously the Langfuse integration only traced in-process HTTP requests (chat / search / eval), so file uploads and every downstream asynq task (document parse, chunk embedding, OCR/VLM, summary / question gen, wiki ingest, datasource sync, etc.) produced either disconnected shallow traces or no observation at all. This change threads one trace end-to-end: - tracer: add SPAN observation type and StartSpan; add ResumeTrace so a worker can attach to an upstream trace without emitting a duplicate trace-create; StartGeneration now auto-picks parentObservationId from ctx so nested trace -> span -> generation trees render correctly. - types.TracingContext + LangfuseTracingCarrier: embed on all 17 asynq payloads so trace_id / parent_obs_id / user_id / session_id serialise into every job. - langfuse.InjectTracing: injected at 28 enqueue sites before json.Marshal so the HTTP-layer trace survives the Redis hop. - langfuse.AsynqMiddleware: mux.Use hook that peeks the payload, either resumes the upstream trace or opens a standalone asynq.<type> trace for scheduled jobs, and wraps the handler in a SPAN with task metadata (id / queue / retry / payload_bytes) plus ERROR level on failure. - GinMiddleware.shouldTrace: whitelist ingestion / knowledge-mutation / FAQ / wiki / datasource endpoints so the root trace actually starts. - Tests: tracer_test.go covers span nesting, error status, and ResumeTrace no-trace-create guarantee; asynq_test.go covers InjectTracing round-trip, middleware resume path, and standalone trace fallback. - Docs: docs/Langfuse\u96c6\u6210.md now lists the covered task types and documents the cross-process propagation model. No behavioural change when Langfuse is disabled (all new code paths are no-ops and carriers serialise to empty strings with omitempty).
68 lines
2.3 KiB
Go
68 lines
2.3 KiB
Go
package langfuse
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|
)
|
|
|
|
// traceCtxKey is the exported context key defined in types/const.go. It lives
|
|
// there (not inside this package) so that logger.CloneContext — which rebuilds
|
|
// a stripped-down context on every request — can preserve the Langfuse trace
|
|
// without importing this package. If we kept the key private here, every
|
|
// CloneContext call would drop the trace and downstream LLM wrappers would
|
|
// each auto-create their own shallow trace, fragmenting a single HTTP request
|
|
// into many unrelated traces in the Langfuse UI.
|
|
var traceCtxKey = types.LangfuseTraceContextKey
|
|
|
|
// parentObsCtxKey tracks the current "parent observation" id — i.e. the span
|
|
// that encloses any generation/span started under it. Unlike traceCtxKey, this
|
|
// key is private: spans are a pure tracer concern and logger.CloneContext
|
|
// should reset the parent across request boundaries (we don't want a leftover
|
|
// span id from one request leaking into another's LLM calls).
|
|
type parentObsCtxKeyType struct{}
|
|
|
|
var parentObsCtxKey = parentObsCtxKeyType{}
|
|
|
|
// withTrace stores a *Trace on the context so downstream LLM wrappers can
|
|
// attach their generations to it.
|
|
func withTrace(ctx context.Context, t *Trace) context.Context {
|
|
if t == nil || ctx == nil {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, traceCtxKey, t)
|
|
}
|
|
|
|
// traceFromCtx retrieves the active trace, if any.
|
|
func traceFromCtx(ctx context.Context) (*Trace, bool) {
|
|
if ctx == nil {
|
|
return nil, false
|
|
}
|
|
t, ok := ctx.Value(traceCtxKey).(*Trace)
|
|
return t, ok && t != nil
|
|
}
|
|
|
|
// TraceFromContext is the public accessor used by HTTP middlewares and
|
|
// handlers that want to set the trace input/output on the active trace.
|
|
func TraceFromContext(ctx context.Context) (*Trace, bool) {
|
|
return traceFromCtx(ctx)
|
|
}
|
|
|
|
// withParentObservation stores the id of the enclosing span so children
|
|
// (sub-spans or generations) can attach via parentObservationId.
|
|
func withParentObservation(ctx context.Context, id string) context.Context {
|
|
if ctx == nil || id == "" {
|
|
return ctx
|
|
}
|
|
return context.WithValue(ctx, parentObsCtxKey, id)
|
|
}
|
|
|
|
// parentObservationFromCtx returns the enclosing span id if any.
|
|
func parentObservationFromCtx(ctx context.Context) (string, bool) {
|
|
if ctx == nil {
|
|
return "", false
|
|
}
|
|
v, ok := ctx.Value(parentObsCtxKey).(string)
|
|
return v, ok && v != ""
|
|
}
|