Skip to content

Commit cfef652

Browse files
richiejprichiejp
committed
feat(api): Allow tracing of requests and responses
Signed-off-by: Richard Palethorpe <[email protected]>
1 parent 878c9d4 commit cfef652

File tree

5 files changed

+186
-2
lines changed

5 files changed

+186
-2
lines changed

‎core/cli/run.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ type RunCMD struct {
7878
DisableGalleryEndpoint bool `env:"LOCALAI_DISABLE_GALLERY_ENDPOINT,DISABLE_GALLERY_ENDPOINT" help:"Disable the gallery endpoints" group:"api"`
7979
MachineTag string `env:"LOCALAI_MACHINE_TAG,MACHINE_TAG" help:"Add Machine-Tag header to each response which is useful to track the machine in the P2P network" group:"api"`
8080
LoadToMemory []string `env:"LOCALAI_LOAD_TO_MEMORY,LOAD_TO_MEMORY" help:"A list of models to load into memory at startup" group:"models"`
81+
EnableTracing bool `env:"LOCALAI_ENABLE_TRACING,ENABLE_TRACING" help:"Enable API tracing" group:"api"`
8182
AgentJobRetentionDays int `env:"LOCALAI_AGENT_JOB_RETENTION_DAYS,AGENT_JOB_RETENTION_DAYS" default:"30" help:"Number of days to keep agent job history (default: 30)" group:"api"`
8283

8384
Version bool
@@ -150,6 +151,10 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
150151
opts = append(opts, config.DisableRuntimeSettings)
151152
}
152153

154+
if r.EnableTracing {
155+
opts = append(opts, config.EnableTracing)
156+
}
157+
153158
token := ""
154159
if r.Peer2Peer || r.Peer2PeerToken != "" {
155160
log.Info().Msg("P2P mode enabled")

‎core/config/application_config.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ type ApplicationConfig struct {
1919
UploadLimitMB, Threads, ContextSize int
2020
F16 bool
2121
Debug bool
22+
EnableTracing bool
2223
GeneratedContentDir string
2324

2425
UploadDir string
@@ -158,6 +159,10 @@ var EnableWatchDog = func(o *ApplicationConfig) {
158159
o.WatchDog = true
159160
}
160161

162+
var EnableTracing = func(o *ApplicationConfig) {
163+
o.EnableTracing = true
164+
}
165+
161166
var EnableWatchDogIdleCheck = func(o *ApplicationConfig) {
162167
o.WatchDog = true
163168
o.WatchDogIdle = true

‎core/http/middleware/trace.go‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
package middleware
2+
3+
import (
4+
"bytes"
5+
"io"
6+
"net/http"
7+
"sync"
8+
9+
"github.com/labstack/echo/v4"
10+
"github.com/mudler/LocalAI/core/application"
11+
"github.com/rs/zerolog/log"
12+
)
13+
14+
type APIExchange struct {
15+
Request struct {
16+
Method string
17+
Path string
18+
Headers http.Header
19+
Body []byte
20+
}
21+
Response struct {
22+
Status int
23+
Headers http.Header
24+
Body []byte
25+
}
26+
}
27+
28+
var apiLogs []APIExchange
29+
var mu sync.Mutex
30+
var logChan = make(chan APIExchange, 100) // Buffered channel for serialization
31+
32+
func init() {
33+
go func() {
34+
for exchange := range logChan {
35+
mu.Lock()
36+
apiLogs = append(apiLogs, exchange)
37+
mu.Unlock()
38+
log.Debug().Msgf("Logged exchange: %s %s - Status: %d", exchange.Request.Method, exchange.Request.Path, exchange.Response.Status)
39+
}
40+
}()
41+
}
42+
43+
type bodyWriter struct {
44+
http.ResponseWriter
45+
body *bytes.Buffer
46+
}
47+
48+
func (w *bodyWriter) Write(b []byte) (int, error) {
49+
w.body.Write(b)
50+
return w.ResponseWriter.Write(b)
51+
}
52+
53+
func (w *bodyWriter) Flush() {
54+
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
55+
flusher.Flush()
56+
}
57+
}
58+
59+
// TraceMiddleware intercepts and logs JSON API requests and responses
60+
func TraceMiddleware(app *application.Application) echo.MiddlewareFunc {
61+
return func(next echo.HandlerFunc) echo.HandlerFunc {
62+
return func(c echo.Context) error {
63+
if !app.ApplicationConfig().EnableTracing {
64+
return next(c)
65+
}
66+
67+
// Only log if Content-Type is application/json
68+
if c.Request().Header.Get("Content-Type") != "application/json" {
69+
return next(c)
70+
}
71+
72+
body, err := io.ReadAll(c.Request().Body)
73+
if err != nil {
74+
log.Error().Err(err).Msg("Failed to read request body")
75+
return err
76+
}
77+
78+
// Restore the body for downstream handlers
79+
c.Request().Body = io.NopCloser(bytes.NewBuffer(body))
80+
81+
// Wrap response writer to capture body
82+
resBody := new(bytes.Buffer)
83+
mw := &bodyWriter{
84+
ResponseWriter: c.Response().Writer,
85+
body: resBody,
86+
}
87+
c.Response().Writer = mw
88+
89+
err = next(c)
90+
if err != nil {
91+
c.Response().Writer = mw.ResponseWriter // Restore original writer if error
92+
return err
93+
}
94+
95+
// Create exchange log
96+
exchange := APIExchange{
97+
Request: struct {
98+
Method string
99+
Path string
100+
Headers http.Header
101+
Body []byte
102+
}{
103+
Method: c.Request().Method,
104+
Path: c.Path(),
105+
Headers: c.Request().Header.Clone(),
106+
Body: body,
107+
},
108+
Response: struct {
109+
Status int
110+
Headers http.Header
111+
Body []byte
112+
}{
113+
Status: c.Response().Status,
114+
Headers: c.Response().Header().Clone(),
115+
Body: resBody.Bytes(),
116+
},
117+
}
118+
119+
// Send to channel (non-blocking)
120+
select {
121+
case logChan <- exchange:
122+
default:
123+
log.Warn().Msg("API log channel full, dropping log")
124+
}
125+
126+
return nil
127+
}
128+
}
129+
}
130+
131+
// GetAPILogs returns a copy of the logged API exchanges for display
132+
func GetAPILogs() []APIExchange {
133+
mu.Lock()
134+
defer mu.Unlock()
135+
return append([]APIExchange{}, apiLogs...)
136+
}
137+
138+
// ClearAPILogs clears the in-memory logs
139+
func ClearAPILogs() {
140+
mu.Lock()
141+
apiLogs = nil
142+
mu.Unlock()
143+
}

‎core/http/routes/openai.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,18 @@ func RegisterOpenAIRoutes(app *echo.Echo,
1414
re *middleware.RequestExtractor,
1515
application *application.Application) {
1616
// openAI compatible API endpoint
17+
traceMiddleware := middleware.TraceMiddleware(application)
1718

1819
// realtime
1920
// TODO: Modify/disable the API key middleware for this endpoint to allow ephemeral keys created by sessions
2021
app.GET("/v1/realtime", openai.Realtime(application))
21-
app.POST("/v1/realtime/sessions", openai.RealtimeTranscriptionSession(application))
22-
app.POST("/v1/realtime/transcription_session", openai.RealtimeTranscriptionSession(application))
22+
app.POST("/v1/realtime/sessions", openai.RealtimeTranscriptionSession(application), traceMiddleware)
23+
app.POST("/v1/realtime/transcription_session", openai.RealtimeTranscriptionSession(application), traceMiddleware)
2324

2425
// chat
2526
chatHandler := openai.ChatEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig())
2627
chatMiddleware := []echo.MiddlewareFunc{
28+
traceMiddleware,
2729
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)),
2830
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
2931
func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -41,6 +43,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
4143
// edit
4244
editHandler := openai.EditEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig())
4345
editMiddleware := []echo.MiddlewareFunc{
46+
traceMiddleware,
4447
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_EDIT)),
4548
re.BuildConstantDefaultModelNameMiddleware("gpt-4o"),
4649
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
@@ -59,6 +62,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
5962
// completion
6063
completionHandler := openai.CompletionEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig())
6164
completionMiddleware := []echo.MiddlewareFunc{
65+
traceMiddleware,
6266
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_COMPLETION)),
6367
re.BuildConstantDefaultModelNameMiddleware("gpt-4o"),
6468
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
@@ -78,6 +82,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
7882
// MCPcompletion
7983
mcpCompletionHandler := openai.MCPCompletionEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig())
8084
mcpCompletionMiddleware := []echo.MiddlewareFunc{
85+
traceMiddleware,
8186
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)),
8287
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
8388
func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -95,6 +100,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
95100
// embeddings
96101
embeddingHandler := openai.EmbeddingsEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig())
97102
embeddingMiddleware := []echo.MiddlewareFunc{
103+
traceMiddleware,
98104
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_EMBEDDINGS)),
99105
re.BuildConstantDefaultModelNameMiddleware("gpt-4o"),
100106
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
@@ -113,6 +119,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
113119

114120
audioHandler := openai.TranscriptEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig())
115121
audioMiddleware := []echo.MiddlewareFunc{
122+
traceMiddleware,
116123
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_TRANSCRIPT)),
117124
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
118125
func(next echo.HandlerFunc) echo.HandlerFunc {
@@ -130,6 +137,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
130137

131138
audioSpeechHandler := localai.TTSEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig())
132139
audioSpeechMiddleware := []echo.MiddlewareFunc{
140+
traceMiddleware,
133141
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_TTS)),
134142
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.TTSRequest) }),
135143
}
@@ -140,6 +148,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
140148
// images
141149
imageHandler := openai.ImageEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig())
142150
imageMiddleware := []echo.MiddlewareFunc{
151+
traceMiddleware,
143152
// Default: use the first available image generation model
144153
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_IMAGE)),
145154
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
@@ -164,6 +173,7 @@ func RegisterOpenAIRoutes(app *echo.Echo,
164173
// videos (OpenAI-compatible endpoints mapped to LocalAI video handler)
165174
videoHandler := openai.VideoEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig())
166175
videoMiddleware := []echo.MiddlewareFunc{
176+
traceMiddleware,
167177
re.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_VIDEO)),
168178
re.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
169179
func(next echo.HandlerFunc) echo.HandlerFunc {

‎core/http/routes/ui_api.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/mudler/LocalAI/core/config"
1717
"github.com/mudler/LocalAI/core/gallery"
1818
"github.com/mudler/LocalAI/core/http/endpoints/localai"
19+
"github.com/mudler/LocalAI/core/http/middleware"
1920
"github.com/mudler/LocalAI/core/p2p"
2021
"github.com/mudler/LocalAI/core/services"
2122
"github.com/mudler/LocalAI/pkg/model"
@@ -947,4 +948,24 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model
947948
app.GET("/api/settings", localai.GetSettingsEndpoint(applicationInstance))
948949
app.POST("/api/settings", localai.UpdateSettingsEndpoint(applicationInstance))
949950
}
951+
952+
// Logs API
953+
app.GET("/api/traces", func(c echo.Context) error {
954+
if !appConfig.EnableTracing {
955+
return c.JSON(503, map[string]any{
956+
"error": "Tracing disabled",
957+
})
958+
}
959+
logs := middleware.GetAPILogs()
960+
return c.JSON(200, map[string]interface{}{
961+
"logs": logs,
962+
})
963+
})
964+
965+
app.POST("/api/traces/clear", func(c echo.Context) error {
966+
middleware.ClearAPILogs()
967+
return c.JSON(200, map[string]interface{}{
968+
"message": "Traces cleared",
969+
})
970+
})
950971
}

0 commit comments

Comments
 (0)