Spaces:
Sleeping
Sleeping
| // Package server is the HTTP API in front of the memory store. | |
| package server | |
| import ( | |
| "crypto/subtle" | |
| "encoding/json" | |
| "errors" | |
| "fmt" | |
| "log/slog" | |
| "net/http" | |
| "strconv" | |
| "sync/atomic" | |
| "time" | |
| "github.com/shaikn6/lethe/internal/memory" | |
| "github.com/shaikn6/lethe/internal/policy" | |
| ) | |
| const ( | |
| maxBodyBytes = 1 << 20 | |
| defaultK = 5 | |
| ) | |
| type rememberRequest struct { | |
| Namespace string `json:"namespace"` | |
| Text string `json:"text"` | |
| Embedding []float32 `json:"embedding"` | |
| Importance *float64 `json:"importance"` | |
| CreatedAt time.Time `json:"created_at"` | |
| Metadata map[string]string `json:"metadata"` | |
| } | |
| type recallRequest struct { | |
| Namespace string `json:"namespace"` | |
| Query string `json:"query"` | |
| Embedding []float32 `json:"embedding"` | |
| K int `json:"k"` | |
| AsOf time.Time `json:"as_of"` | |
| Explore bool `json:"explore"` | |
| } | |
| type policyRequest struct { | |
| Weights []float64 `json:"weights"` | |
| } | |
| type feedbackRequest struct { | |
| RecallID string `json:"recall_id"` | |
| Reward *float64 `json:"reward"` | |
| Useful []string `json:"useful"` // optional: ids of the returned memories that helped | |
| } | |
| // Config holds the settings of the HTTP API. | |
| type Config struct { | |
| // Token protects the memory endpoints. Empty leaves them open, which is | |
| // only reasonable on a loopback address. | |
| Token string | |
| // AdminToken protects the endpoints that change how the server behaves | |
| // for everyone: setting the ranking weights and forcing a compaction. | |
| // Empty switches those endpoints off. | |
| AdminToken string | |
| // MaxInFlight is how many requests are served at once before the rest | |
| // are turned away with 503. Default 256. | |
| MaxInFlight int | |
| // CompactEvery is the shortest gap between compactions requested over | |
| // the API, because each one briefly holds every other request. Default 30s. | |
| CompactEvery time.Duration | |
| // EventCounters, when set, is reported by the health check so that a | |
| // stalled event log is visible. | |
| EventCounters func() any | |
| } | |
| // New returns the API handler. | |
| func New(store *memory.Store, cfg Config, log *slog.Logger) http.Handler { | |
| if cfg.MaxInFlight <= 0 { | |
| cfg.MaxInFlight = 256 | |
| } | |
| if cfg.CompactEvery <= 0 { | |
| cfg.CompactEvery = 30 * time.Second | |
| } | |
| mux := http.NewServeMux() | |
| // Unhealthy means the journal can no longer be written: the server still | |
| // answers recalls but refuses writes, so whatever runs it should replace it. | |
| mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { | |
| stats := store.Stats() | |
| status := http.StatusOK | |
| if !stats.JournalHealthy { | |
| status = http.StatusServiceUnavailable | |
| } | |
| body := struct { | |
| memory.Stats | |
| Events any `json:"events,omitempty"` | |
| }{Stats: stats} | |
| if cfg.EventCounters != nil { | |
| body.Events = cfg.EventCounters() | |
| } | |
| writeJSON(w, status, body) | |
| }) | |
| mux.HandleFunc("POST /v1/memories", func(w http.ResponseWriter, r *http.Request) { | |
| var req rememberRequest | |
| if !decode(w, r, &req) { | |
| return | |
| } | |
| importance := 0.5 | |
| if req.Importance != nil { | |
| importance = *req.Importance | |
| } | |
| m, err := store.Remember(memory.Memory{ | |
| Namespace: req.Namespace, Text: req.Text, Embedding: req.Embedding, | |
| Importance: importance, CreatedAt: req.CreatedAt, Metadata: req.Metadata, | |
| }) | |
| if err != nil { | |
| writeError(w, log, err) | |
| return | |
| } | |
| m.Embedding = nil | |
| writeJSON(w, http.StatusCreated, m) | |
| }) | |
| mux.HandleFunc("DELETE /v1/memories/{namespace}/{id}", func(w http.ResponseWriter, r *http.Request) { | |
| if err := store.Forget(r.PathValue("namespace"), r.PathValue("id")); err != nil { | |
| writeError(w, log, err) | |
| return | |
| } | |
| w.WriteHeader(http.StatusNoContent) | |
| }) | |
| mux.HandleFunc("POST /v1/recall", func(w http.ResponseWriter, r *http.Request) { | |
| var req recallRequest | |
| if !decode(w, r, &req) { | |
| return | |
| } | |
| if req.K == 0 { | |
| req.K = defaultK | |
| } | |
| out, err := store.Recall(memory.Query{ | |
| Namespace: req.Namespace, Text: req.Query, Embedding: req.Embedding, | |
| K: req.K, AsOf: req.AsOf, Explore: req.Explore, | |
| }) | |
| if err != nil { | |
| writeError(w, log, err) | |
| return | |
| } | |
| writeJSON(w, http.StatusOK, out) | |
| }) | |
| mux.HandleFunc("POST /v1/feedback", func(w http.ResponseWriter, r *http.Request) { | |
| var req feedbackRequest | |
| if !decode(w, r, &req) { | |
| return | |
| } | |
| if req.Reward == nil { | |
| http.Error(w, "reward is required", http.StatusBadRequest) | |
| return | |
| } | |
| if len(req.Useful) > memory.MaxK { | |
| http.Error(w, "too many useful ids", http.StatusBadRequest) | |
| return | |
| } | |
| if err := store.Feedback(req.RecallID, *req.Reward, req.Useful); err != nil { | |
| writeError(w, log, err) | |
| return | |
| } | |
| w.WriteHeader(http.StatusNoContent) | |
| }) | |
| root := http.NewServeMux() | |
| root.Handle("/", authorize(cfg.Token, mux)) | |
| if cfg.AdminToken != "" { | |
| admin := http.NewServeMux() | |
| var lastCompact atomic.Int64 | |
| admin.HandleFunc("POST /v1/compact", func(w http.ResponseWriter, _ *http.Request) { | |
| now, last := time.Now().UnixNano(), lastCompact.Load() | |
| if wait := cfg.CompactEvery - time.Duration(now-last); last != 0 && wait > 0 { | |
| w.Header().Set("Retry-After", strconv.Itoa(int(wait.Seconds())+1)) | |
| http.Error(w, "compaction was requested moments ago", http.StatusTooManyRequests) | |
| return | |
| } | |
| if !lastCompact.CompareAndSwap(last, now) { | |
| http.Error(w, "compaction already requested", http.StatusTooManyRequests) | |
| return | |
| } | |
| if err := store.Compact(); err != nil { | |
| writeError(w, log, err) | |
| return | |
| } | |
| w.WriteHeader(http.StatusNoContent) | |
| }) | |
| admin.HandleFunc("PUT /v1/policy", func(w http.ResponseWriter, r *http.Request) { | |
| var req policyRequest | |
| if !decode(w, r, &req) { | |
| return | |
| } | |
| // A fixed-size array would silently drop or zero-fill a list of the | |
| // wrong length, so the length is checked on a slice first. | |
| if len(req.Weights) != policy.NumFeatures { | |
| http.Error(w, fmt.Sprintf("weights must have exactly %d values", policy.NumFeatures), http.StatusBadRequest) | |
| return | |
| } | |
| if err := store.SetWeights(policy.Features(req.Weights)); err != nil { | |
| writeError(w, log, err) | |
| return | |
| } | |
| w.WriteHeader(http.StatusNoContent) | |
| }) | |
| protected := authorize(cfg.AdminToken, admin) | |
| root.Handle("POST /v1/compact", protected) | |
| root.Handle("PUT /v1/policy", protected) | |
| } | |
| return recoverPanics(log, limitInFlight(cfg.MaxInFlight, root)) | |
| } | |
| // limitInFlight sheds load instead of letting requests pile up without | |
| // bound. Health checks always get through. | |
| func limitInFlight(n int, next http.Handler) http.Handler { | |
| slots := make(chan struct{}, n) | |
| return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| if r.URL.Path == "/healthz" { | |
| next.ServeHTTP(w, r) | |
| return | |
| } | |
| select { | |
| case slots <- struct{}{}: | |
| defer func() { <-slots }() | |
| next.ServeHTTP(w, r) | |
| default: | |
| w.Header().Set("Retry-After", "1") | |
| http.Error(w, "server is at capacity", http.StatusServiceUnavailable) | |
| } | |
| }) | |
| } | |
| func authorize(token string, next http.Handler) http.Handler { | |
| if token == "" { | |
| return next | |
| } | |
| want := []byte("Bearer " + token) | |
| return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| got := []byte(r.Header.Get("Authorization")) | |
| if r.URL.Path != "/healthz" && subtle.ConstantTimeCompare(got, want) != 1 { | |
| http.Error(w, "unauthorized", http.StatusUnauthorized) | |
| return | |
| } | |
| next.ServeHTTP(w, r) | |
| }) | |
| } | |
| // recoverPanics keeps one failing request from taking the process down. | |
| func recoverPanics(log *slog.Logger, next http.Handler) http.Handler { | |
| return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| defer func() { | |
| if v := recover(); v != nil { | |
| log.Error("panic in handler", "path", r.URL.Path, "panic", v) | |
| http.Error(w, "internal error", http.StatusInternalServerError) | |
| } | |
| }() | |
| next.ServeHTTP(w, r) | |
| }) | |
| } | |
| func decode(w http.ResponseWriter, r *http.Request, into any) bool { | |
| dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, maxBodyBytes)) | |
| dec.DisallowUnknownFields() | |
| if err := dec.Decode(into); err != nil { | |
| http.Error(w, "invalid JSON body", http.StatusBadRequest) | |
| return false | |
| } | |
| return true | |
| } | |
| func writeError(w http.ResponseWriter, log *slog.Logger, err error) { | |
| switch { | |
| case errors.Is(err, memory.ErrInvalid): | |
| http.Error(w, err.Error(), http.StatusBadRequest) | |
| case errors.Is(err, memory.ErrNotFound): | |
| http.Error(w, "not found", http.StatusNotFound) | |
| case errors.Is(err, memory.ErrClosed): | |
| http.Error(w, "server is shutting down", http.StatusServiceUnavailable) | |
| case errors.Is(err, memory.ErrFull): | |
| http.Error(w, "store is full", http.StatusInsufficientStorage) | |
| default: | |
| log.Error("request failed", "error", err) | |
| http.Error(w, "internal error", http.StatusInternalServerError) | |
| } | |
| } | |
| func writeJSON(w http.ResponseWriter, status int, v any) { | |
| w.Header().Set("Content-Type", "application/json") | |
| w.WriteHeader(status) | |
| _ = json.NewEncoder(w).Encode(v) | |
| } | |