feat(traces): Add backend traces

core/backend/embeddings.go

@@ -2,8 +2,10 @@ package backend
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 
 	"github.com/mudler/LocalAI/pkg/grpc"
 	model "github.com/mudler/LocalAI/pkg/model"
@@ -53,7 +55,7 @@ func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, modelConf
 		}
 	}
 
-	return func() ([]float32, error) {
+	wrappedFn := func() ([]float32, error) {
 		embeds, err := fn()
 		if err != nil {
 			return embeds, err
@@ -67,5 +69,48 @@ func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, modelConf
 			}
 		}
 		return embeds, nil
-	}, nil
+	}
+
+	if appConfig.EnableTracing {
+		trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems)
+
+		traceData := map[string]any{
+			"input_text":         trace.TruncateString(s, 1000),
+			"input_tokens_count": len(tokens),
+		}
+
+		startTime := time.Now()
+		originalFn := wrappedFn
+		wrappedFn = func() ([]float32, error) {
+			result, err := originalFn()
+			duration := time.Since(startTime)
+
+			traceData["embedding_dimensions"] = len(result)
+
+			errStr := ""
+			if err != nil {
+				errStr = err.Error()
+			}
+
+			summary := trace.TruncateString(s, 200)
+			if summary == "" {
+				summary = fmt.Sprintf("tokens[%d]", len(tokens))
+			}
+
+			trace.RecordBackendTrace(trace.BackendTrace{
+				Timestamp: startTime,
+				Duration:  duration,
+				Type:      trace.BackendTraceEmbedding,
+				ModelName: modelConfig.Name,
+				Backend:   modelConfig.Backend,
+				Summary:   summary,
+				Error:     errStr,
+				Data:      traceData,
+			})
+
+			return result, err
+		}
+	}
+
+	return wrappedFn, nil
 }

core/backend/image.go

@@ -1,7 +1,10 @@
 package backend
 
 import (
+	"time"
+
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 
 	"github.com/mudler/LocalAI/pkg/grpc/proto"
 	model "github.com/mudler/LocalAI/pkg/model"
@@ -36,6 +39,46 @@ func ImageGeneration(height, width, step, seed int, positive_prompt, negative_pr
 		return err
 	}
 
+	if appConfig.EnableTracing {
+		trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems)
+
+		traceData := map[string]any{
+			"positive_prompt": positive_prompt,
+			"negative_prompt": negative_prompt,
+			"height":          height,
+			"width":           width,
+			"step":            step,
+			"seed":            seed,
+			"source_image":    src,
+			"destination":     dst,
+		}
+
+		startTime := time.Now()
+		originalFn := fn
+		fn = func() error {
+			err := originalFn()
+			duration := time.Since(startTime)
+
+			errStr := ""
+			if err != nil {
+				errStr = err.Error()
+			}
+
+			trace.RecordBackendTrace(trace.BackendTrace{
+				Timestamp: startTime,
+				Duration:  duration,
+				Type:      trace.BackendTraceImageGeneration,
+				ModelName: modelConfig.Name,
+				Backend:   modelConfig.Backend,
+				Summary:   trace.TruncateString(positive_prompt, 200),
+				Error:     errStr,
+				Data:      traceData,
+			})
+
+			return err
+		}
+	}
+
 	return fn, nil
 }
 

core/backend/llm.go

@@ -7,11 +7,13 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"time"
 	"unicode/utf8"
 
 	"github.com/mudler/xlog"
 
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 	"github.com/mudler/LocalAI/core/schema"
 	"github.com/mudler/LocalAI/core/services"
 
@@ -220,6 +222,84 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima
 		}
 	}
 
+	if o.EnableTracing {
+		trace.InitBackendTracingIfEnabled(o.TracingMaxItems)
+
+		traceData := map[string]any{
+			"prompt":                 s,
+			"use_tokenizer_template": c.TemplateConfig.UseTokenizerTemplate,
+			"chat_template":          c.TemplateConfig.Chat,
+			"function_template":      c.TemplateConfig.Functions,
+			"grammar":               c.Grammar,
+			"stop_words":            c.StopWords,
+			"streaming":             tokenCallback != nil,
+			"images_count":          len(images),
+			"videos_count":          len(videos),
+			"audios_count":          len(audios),
+		}
+
+		if len(messages) > 0 {
+			if msgJSON, err := json.Marshal(messages); err == nil {
+				traceData["messages"] = string(msgJSON)
+			}
+		}
+		if tools != "" {
+			traceData["tools"] = tools
+		}
+		if toolChoice != "" {
+			traceData["tool_choice"] = toolChoice
+		}
+		if reasoningJSON, err := json.Marshal(c.ReasoningConfig); err == nil {
+			traceData["reasoning_config"] = string(reasoningJSON)
+		}
+		traceData["functions_config"] = map[string]any{
+			"grammar_disabled":  c.FunctionsConfig.GrammarConfig.NoGrammar,
+			"parallel_calls":    c.FunctionsConfig.GrammarConfig.ParallelCalls,
+			"mixed_mode":        c.FunctionsConfig.GrammarConfig.MixedMode,
+			"xml_format_preset": c.FunctionsConfig.XMLFormatPreset,
+		}
+		if c.Temperature != nil {
+			traceData["temperature"] = *c.Temperature
+		}
+		if c.TopP != nil {
+			traceData["top_p"] = *c.TopP
+		}
+		if c.Maxtokens != nil {
+			traceData["max_tokens"] = *c.Maxtokens
+		}
+
+		startTime := time.Now()
+		originalFn := fn
+		fn = func() (LLMResponse, error) {
+			resp, err := originalFn()
+			duration := time.Since(startTime)
+
+			traceData["response"] = resp.Response
+			traceData["token_usage"] = map[string]any{
+				"prompt":     resp.Usage.Prompt,
+				"completion": resp.Usage.Completion,
+			}
+
+			errStr := ""
+			if err != nil {
+				errStr = err.Error()
+			}
+
+			trace.RecordBackendTrace(trace.BackendTrace{
+				Timestamp: startTime,
+				Duration:  duration,
+				Type:      trace.BackendTraceLLM,
+				ModelName: c.Name,
+				Backend:   c.Backend,
+				Summary:   trace.GenerateLLMSummary(messages, s),
+				Error:     errStr,
+				Data:      traceData,
+			})
+
+			return resp, err
+		}
+	}
+
 	return fn, nil
 }
 

core/backend/rerank.go

@@ -3,8 +3,10 @@ package backend
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 	"github.com/mudler/LocalAI/pkg/grpc/proto"
 	model "github.com/mudler/LocalAI/pkg/model"
 )
@@ -20,7 +22,35 @@ func Rerank(request *proto.RerankRequest, loader *model.ModelLoader, appConfig *
 		return nil, fmt.Errorf("could not load rerank model")
 	}
 
+	var startTime time.Time
+	if appConfig.EnableTracing {
+		trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems)
+		startTime = time.Now()
+	}
+
 	res, err := rerankModel.Rerank(context.Background(), request)
 
+	if appConfig.EnableTracing {
+		errStr := ""
+		if err != nil {
+			errStr = err.Error()
+		}
+
+		trace.RecordBackendTrace(trace.BackendTrace{
+			Timestamp: startTime,
+			Duration:  time.Since(startTime),
+			Type:      trace.BackendTraceRerank,
+			ModelName: modelConfig.Name,
+			Backend:   modelConfig.Backend,
+			Summary:   trace.TruncateString(request.Query, 200),
+			Error:     errStr,
+			Data: map[string]any{
+				"query":           request.Query,
+				"documents_count": len(request.Documents),
+				"top_n":           request.TopN,
+			},
+		})
+	}
+
 	return res, err
 }

core/backend/soundgeneration.go

@@ -5,8 +5,10 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 	"github.com/mudler/LocalAI/pkg/grpc/proto"
 	"github.com/mudler/LocalAI/pkg/model"
 	"github.com/mudler/LocalAI/pkg/utils"
@@ -92,7 +94,51 @@ func SoundGeneration(
 		req.Instrumental = instrumental
 	}
 
+	var startTime time.Time
+	if appConfig.EnableTracing {
+		trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems)
+		startTime = time.Now()
+	}
+
 	res, err := soundGenModel.SoundGeneration(context.Background(), req)
+
+	if appConfig.EnableTracing {
+		errStr := ""
+		if err != nil {
+			errStr = err.Error()
+		} else if res != nil && !res.Success {
+			errStr = fmt.Sprintf("sound generation error: %s", res.Message)
+		}
+
+		summary := trace.TruncateString(text, 200)
+		if summary == "" && caption != "" {
+			summary = trace.TruncateString(caption, 200)
+		}
+
+		traceData := map[string]any{
+			"text":    text,
+			"caption": caption,
+			"lyrics":  lyrics,
+		}
+		if duration != nil {
+			traceData["duration"] = *duration
+		}
+		if temperature != nil {
+			traceData["temperature"] = *temperature
+		}
+
+		trace.RecordBackendTrace(trace.BackendTrace{
+			Timestamp: startTime,
+			Duration:  time.Since(startTime),
+			Type:      trace.BackendTraceSoundGeneration,
+			ModelName: modelConfig.Name,
+			Backend:   modelConfig.Backend,
+			Summary:   summary,
+			Error:     errStr,
+			Data:      traceData,
+		})
+	}
+
 	if err != nil {
 		return "", nil, err
 	}

core/backend/tokenize.go

@@ -1,7 +1,10 @@
 package backend
 
 import (
+	"time"
+
 	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/trace"
 	"github.com/mudler/LocalAI/core/schema"
 	"github.com/mudler/LocalAI/pkg/grpc"
 	"github.com/mudler/LocalAI/pkg/model"
@@ -21,8 +24,41 @@ func ModelTokenize(s string, loader *model.ModelLoader, modelConfig config.Model
 	predictOptions := gRPCPredictOpts(modelConfig, loader.ModelPath)
 	predictOptions.Prompt = s
 
+	var startTime time.Time
+	if appConfig.EnableTracing {
+		trace.InitBackendTracingIfEnabled(appConfig.TracingMaxItems)
+		startTime = time.Now()
+	}
+
 	// tokenize the string
 	resp, err := inferenceModel.TokenizeString(appConfig.Context, predictOptions)
+
+	if appConfig.EnableTracing {
+		errStr := ""
+		if err != nil {
+			errStr = err.Error()
+		}
+
+		tokenCount := 0
+		if resp.Tokens != nil {
+			tokenCount = len(resp.Tokens)
+		}
+
+		trace.RecordBackendTrace(trace.BackendTrace{
+			Timestamp: startTime,
+			Duration:  time.Since(startTime),
+			Type:      trace.BackendTraceTokenize,
+			ModelName: modelConfig.Name,
+			Backend:   modelConfig.Backend,
+			Summary:   trace.TruncateString(s, 200),
+			Error:     errStr,
+			Data: map[string]any{
+				"input_text":  trace.TruncateString(s, 1000),
+				"token_count": tokenCount,
+			},
+		})
+	}
+
 	if err != nil {
 		return schema.TokenizeResponse{}, err
 	}