Fix service shutdown
This commit is contained in:
+148
-383
@@ -6,7 +6,7 @@ package service
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"embed"
|
"embed"
|
||||||
"encoding/json"
|
"encoding/json/v2"
|
||||||
"fmt"
|
"fmt"
|
||||||
"html/template"
|
"html/template"
|
||||||
"io/fs"
|
"io/fs"
|
||||||
@@ -19,13 +19,16 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
"uuid"
|
||||||
|
|
||||||
"code.chimeric.al/chimerical/odidere/internal/config"
|
"code.chimeric.al/chimerical/odidere/internal/config"
|
||||||
|
"code.chimeric.al/chimerical/odidere/internal/conversation"
|
||||||
|
"code.chimeric.al/chimerical/odidere/internal/db"
|
||||||
"code.chimeric.al/chimerical/odidere/internal/job"
|
"code.chimeric.al/chimerical/odidere/internal/job"
|
||||||
"code.chimeric.al/chimerical/odidere/internal/llm"
|
"code.chimeric.al/chimerical/odidere/internal/llm"
|
||||||
"code.chimeric.al/chimerical/odidere/internal/service/templates"
|
"code.chimeric.al/chimerical/odidere/internal/service/templates"
|
||||||
|
"code.chimeric.al/chimerical/odidere/internal/tool"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
|
||||||
"github.com/robfig/cron/v3"
|
"github.com/robfig/cron/v3"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -48,6 +51,9 @@ var static embed.FS
|
|||||||
type Service struct {
|
type Service struct {
|
||||||
cfg *config.Config
|
cfg *config.Config
|
||||||
cron *cron.Cron
|
cron *cron.Cron
|
||||||
|
db *db.SQLite
|
||||||
|
dbr *db.Queries
|
||||||
|
dbw *db.Queries
|
||||||
jobs []*job.Job
|
jobs []*job.Job
|
||||||
llms map[string]*llm.Client
|
llms map[string]*llm.Client
|
||||||
log *slog.Logger
|
log *slog.Logger
|
||||||
@@ -55,20 +61,49 @@ type Service struct {
|
|||||||
mux *http.ServeMux
|
mux *http.ServeMux
|
||||||
server *http.Server
|
server *http.Server
|
||||||
tmpl *template.Template
|
tmpl *template.Template
|
||||||
tools *llm.Registry
|
tools *tool.Registry
|
||||||
|
|
||||||
|
ctx context.Context
|
||||||
|
cancel context.CancelFunc
|
||||||
|
mu sync.Mutex
|
||||||
|
conversations map[string]*conversation.Conversation
|
||||||
}
|
}
|
||||||
|
|
||||||
// New creates a Service from the provided configuration.
|
// New creates a Service from the provided configuration.
|
||||||
// It initializes all clients and the HTTP server.
|
// It initializes all clients and the HTTP server.
|
||||||
func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
func New(
|
||||||
|
cfg *config.Config,
|
||||||
|
log *slog.Logger,
|
||||||
|
sdb *db.SQLite,
|
||||||
|
) (*Service, error) {
|
||||||
|
if cfg == nil {
|
||||||
|
return nil, fmt.Errorf("config is nil")
|
||||||
|
}
|
||||||
|
if err := cfg.Validate(); err != nil {
|
||||||
|
return nil, fmt.Errorf("invalid config: %w", err)
|
||||||
|
}
|
||||||
|
if log == nil {
|
||||||
|
return nil, fmt.Errorf("logger is nil")
|
||||||
|
}
|
||||||
|
if sdb == nil || sdb.Writer == nil || sdb.Reader == nil {
|
||||||
|
return nil, fmt.Errorf("database is nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
var svc = &Service{
|
var svc = &Service{
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
|
db: sdb,
|
||||||
|
dbr: db.New(sdb.Reader),
|
||||||
|
dbw: db.New(sdb.Writer),
|
||||||
log: log,
|
log: log,
|
||||||
mux: http.NewServeMux(),
|
mux: http.NewServeMux(),
|
||||||
|
ctx: ctx,
|
||||||
|
cancel: cancel,
|
||||||
|
conversations: make(map[string]*conversation.Conversation),
|
||||||
}
|
}
|
||||||
|
|
||||||
// Setup tool registry.
|
// Setup tool registry.
|
||||||
registry, err := llm.NewRegistry(cfg.Tools)
|
registry, err := tool.NewRegistry(cfg.Tools)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("load tools: %v", err)
|
return nil, fmt.Errorf("load tools: %v", err)
|
||||||
}
|
}
|
||||||
@@ -80,11 +115,10 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
|||||||
client, err := llm.NewClient(
|
client, err := llm.NewClient(
|
||||||
llm.Config{
|
llm.Config{
|
||||||
Key: pc.Key,
|
Key: pc.Key,
|
||||||
SystemMessage: cfg.SystemMessage,
|
Provider: pc.Name,
|
||||||
Timeout: pc.Timeout,
|
Timeout: pc.Timeout,
|
||||||
URL: pc.URL,
|
URL: pc.URL,
|
||||||
},
|
},
|
||||||
registry,
|
|
||||||
log,
|
log,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -100,6 +134,9 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
|||||||
slog.String("url", pc.URL),
|
slog.String("url", pc.URL),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
if svc.llms[cfg.DefaultModel.Provider] == nil {
|
||||||
|
return nil, fmt.Errorf("unknown default provider")
|
||||||
|
}
|
||||||
|
|
||||||
// Build model allowed list map: provider name -> set of base model IDs.
|
// Build model allowed list map: provider name -> set of base model IDs.
|
||||||
svc.allowedModels = make(
|
svc.allowedModels = make(
|
||||||
@@ -121,19 +158,12 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
|||||||
jobs := make([]*job.Job, 0, len(cfg.Jobs))
|
jobs := make([]*job.Job, 0, len(cfg.Jobs))
|
||||||
for _, jc := range cfg.Jobs {
|
for _, jc := range cfg.Jobs {
|
||||||
// Resolve the provider and model for this job.
|
// Resolve the provider and model for this job.
|
||||||
var provider, modelID string
|
provider, modelID := jc.Model.Provider, jc.Model.Model
|
||||||
if jc.Model.Provider != "" {
|
if provider == "" {
|
||||||
if jc.Model.Model == "" {
|
provider = cfg.DefaultModel.Provider
|
||||||
return nil, fmt.Errorf(
|
|
||||||
"job %q: model.provider set but model.model is empty",
|
|
||||||
jc.Name,
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
provider = jc.Model.Provider
|
if modelID == "" {
|
||||||
modelID = jc.Model.Model
|
modelID = cfg.DefaultModel.Model
|
||||||
} else {
|
|
||||||
provider = svc.cfg.DefaultModel.Provider
|
|
||||||
modelID = svc.cfg.DefaultModel.Model
|
|
||||||
}
|
}
|
||||||
llmc := svc.llms[provider]
|
llmc := svc.llms[provider]
|
||||||
if llmc == nil {
|
if llmc == nil {
|
||||||
@@ -143,7 +173,14 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
|||||||
provider,
|
provider,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
j, err := job.New(jc, log, llmc, modelID)
|
jc.Model = config.ModelConfig{
|
||||||
|
Provider: provider,
|
||||||
|
Model: modelID,
|
||||||
|
}
|
||||||
|
if jc.SystemMessage == "" {
|
||||||
|
jc.SystemMessage = cfg.SystemMessage
|
||||||
|
}
|
||||||
|
j, err := job.New(jc, llmc, log)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf(
|
return nil, fmt.Errorf(
|
||||||
"create job %q: %w", jc.Name, err,
|
"create job %q: %w", jc.Name, err,
|
||||||
@@ -175,9 +212,46 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
|||||||
"/static/", http.FileServer(http.FS(staticFS)),
|
"/static/", http.FileServer(http.FS(staticFS)),
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
svc.mux.HandleFunc("POST /v1/chat/voice", svc.chat)
|
svc.mux.HandleFunc("GET /models", svc.models)
|
||||||
svc.mux.HandleFunc("POST /v1/chat/voice/stream", svc.chatStream)
|
|
||||||
svc.mux.HandleFunc("GET /v1/models", svc.models)
|
svc.mux.HandleFunc("GET /conversations", svc.listConversations)
|
||||||
|
svc.mux.HandleFunc("POST /conversations", svc.createConversation)
|
||||||
|
svc.mux.HandleFunc("GET /conversations/{id}", svc.getConversationHandler)
|
||||||
|
svc.mux.HandleFunc("DELETE /conversations/{id}", svc.deleteConversation)
|
||||||
|
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"POST /conversations/{id}/inference", svc.inference,
|
||||||
|
)
|
||||||
|
svc.mux.HandleFunc("POST /conversations/{id}/stop", svc.stop)
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"POST /conversations/{id}/subscribe", svc.subscribe,
|
||||||
|
)
|
||||||
|
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"POST /conversations/{id}/messages", svc.insertMessage,
|
||||||
|
)
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"PUT /conversations/{id}/messages/{mid}", svc.updateMessage,
|
||||||
|
)
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"DELETE /conversations/{id}/messages/{mid}", svc.deleteMessage,
|
||||||
|
)
|
||||||
|
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"POST /conversations/{id}/messages/{mid}/calls",
|
||||||
|
svc.addToolCall,
|
||||||
|
)
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"PUT /conversations/{id}/messages/{mid}/calls/{cid}",
|
||||||
|
svc.updateToolCall,
|
||||||
|
)
|
||||||
|
svc.mux.HandleFunc(
|
||||||
|
"DELETE /conversations/{id}/messages/{mid}/calls/{cid}",
|
||||||
|
svc.deleteToolCall,
|
||||||
|
)
|
||||||
|
|
||||||
|
svc.mux.HandleFunc("GET /executions", svc.listExecutions)
|
||||||
|
svc.mux.HandleFunc("PUT /executions/{id}", svc.updateExecution)
|
||||||
|
|
||||||
svc.server = &http.Server{
|
svc.server = &http.Server{
|
||||||
Addr: cfg.Address,
|
Addr: cfg.Address,
|
||||||
@@ -192,7 +266,7 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
|||||||
func (svc *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
func (svc *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
var (
|
var (
|
||||||
start = time.Now()
|
start = time.Now()
|
||||||
id = uuid.NewString()
|
id = uuid.NewV4().String()
|
||||||
ip = func() string {
|
ip = func() string {
|
||||||
if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
|
if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
|
||||||
before, _, found := strings.Cut(ip, ",")
|
before, _, found := strings.Cut(ip, ",")
|
||||||
@@ -296,7 +370,7 @@ func (svc *Service) Run(ctx context.Context) error {
|
|||||||
"jobs",
|
"jobs",
|
||||||
slog.Int("count", svc.JobCount()),
|
slog.Int("count", svc.JobCount()),
|
||||||
),
|
),
|
||||||
slog.Any("shutdown_timeout", svc.cfg.GetShutdownTimeout()),
|
slog.Any("shutdown_timeout", svc.cfg.ShutdownTimeout),
|
||||||
)
|
)
|
||||||
|
|
||||||
// Setup signal handling for graceful shutdown.
|
// Setup signal handling for graceful shutdown.
|
||||||
@@ -305,26 +379,29 @@ func (svc *Service) Run(ctx context.Context) error {
|
|||||||
os.Interrupt, syscall.SIGTERM,
|
os.Interrupt, syscall.SIGTERM,
|
||||||
)
|
)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
svc.ctx, svc.cancel = context.WithCancel(ctx)
|
||||||
|
defer svc.cancel()
|
||||||
|
|
||||||
// Register jobs with cron scheduler.
|
// Register jobs with cron scheduler.
|
||||||
svc.cron = cron.New()
|
svc.cron = cron.New()
|
||||||
for _, j := range svc.jobs {
|
for _, j := range svc.jobs {
|
||||||
entry, err := svc.cron.AddFunc(j.Schedule(), func() {
|
entry, err := svc.cron.AddFunc(j.Schedule(), func() {
|
||||||
result := j.Run(ctx)
|
title := j.Name()
|
||||||
if result.Error != nil {
|
system := j.SystemMessage()
|
||||||
|
_, c, err := svc.newConversation(
|
||||||
|
svc.ctx, j.UserId(), &title, system,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
svc.log.Error("failed to create job conversation",
|
||||||
|
slog.String("name", j.Name()),
|
||||||
|
slog.Any("error", err),
|
||||||
|
)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := j.Run(svc.ctx, c); err != nil {
|
||||||
svc.log.Error("job failed",
|
svc.log.Error("job failed",
|
||||||
slog.String("name", j.Name()),
|
slog.String("name", j.Name()),
|
||||||
slog.Any("error", result.Error),
|
slog.Any("error", err),
|
||||||
slog.Duration(
|
|
||||||
"duration", result.Duration,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
svc.log.Info("job completed",
|
|
||||||
slog.String("name", j.Name()),
|
|
||||||
slog.Duration(
|
|
||||||
"duration", result.Duration,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
@@ -358,36 +435,56 @@ func (svc *Service) Run(ctx context.Context) error {
|
|||||||
}()
|
}()
|
||||||
|
|
||||||
// Wait for shutdown signal or server error.
|
// Wait for shutdown signal or server error.
|
||||||
|
var runErr error
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
svc.log.Info("shutdown signal received")
|
svc.log.Info("shutdown signal received")
|
||||||
case err := <-errs:
|
case err := <-errs:
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("server error: %w", err)
|
runErr = fmt.Errorf("server error: %w", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Graceful shutdown with timeout.
|
// Use one deadline for HTTP shutdown and job completion.
|
||||||
|
// The shutdown context must remain active after svc.ctx is canceled.
|
||||||
shutdownCtx, shutdownCancel := context.WithTimeout(
|
shutdownCtx, shutdownCancel := context.WithTimeout(
|
||||||
context.Background(),
|
context.Background(),
|
||||||
svc.cfg.GetShutdownTimeout(),
|
svc.cfg.ShutdownTimeout,
|
||||||
)
|
)
|
||||||
defer shutdownCancel()
|
defer shutdownCancel()
|
||||||
|
|
||||||
// Stop cron scheduler — in-flight jobs continue to completion.
|
svc.cancel()
|
||||||
svc.log.Info("stopping job scheduler")
|
|
||||||
svc.cron.Stop()
|
|
||||||
|
|
||||||
|
svc.log.Info("stopping job scheduler")
|
||||||
|
cronCtx := svc.cron.Stop()
|
||||||
|
|
||||||
|
var shutdownErr error
|
||||||
svc.log.Info("shutting down HTTP server")
|
svc.log.Info("shutting down HTTP server")
|
||||||
if err := svc.server.Shutdown(shutdownCtx); err != nil {
|
if err := svc.server.Shutdown(shutdownCtx); err != nil {
|
||||||
svc.log.Warn(
|
shutdownErr = fmt.Errorf("HTTP shutdown: %w", err)
|
||||||
"shutdown timeout reached",
|
} else {
|
||||||
slog.Any("error", err),
|
// Jobs can save tool results after cancellation. Wait until
|
||||||
|
// their callbacks return before the caller closes the database.
|
||||||
|
select {
|
||||||
|
case <-cronCtx.Done():
|
||||||
|
case <-shutdownCtx.Done():
|
||||||
|
shutdownErr = fmt.Errorf(
|
||||||
|
"wait for jobs: %w", shutdownCtx.Err(),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// On failure, main exits without closing the database because
|
||||||
|
// handlers or jobs may still use it.
|
||||||
|
if shutdownErr != nil {
|
||||||
|
if runErr != nil {
|
||||||
|
return fmt.Errorf("%w; %w", runErr, shutdownErr)
|
||||||
|
}
|
||||||
|
return shutdownErr
|
||||||
|
}
|
||||||
|
|
||||||
svc.log.Info("terminating")
|
svc.log.Info("terminating")
|
||||||
return nil
|
return runErr
|
||||||
}
|
}
|
||||||
|
|
||||||
// ListJobs returns all registered jobs.
|
// ListJobs returns all registered jobs.
|
||||||
@@ -438,338 +535,6 @@ func (svc *Service) status(w http.ResponseWriter, r *http.Request) {
|
|||||||
w.WriteHeader(http.StatusOK)
|
w.WriteHeader(http.StatusOK)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Request is the incoming request format for the chat endpoints.
|
|
||||||
type Request struct {
|
|
||||||
// Messages is the conversation history.
|
|
||||||
Messages []llm.Message `json:"messages"`
|
|
||||||
// Provider is the LLM provider name. If empty, the default provider is used.
|
|
||||||
Provider string `json:"provider,omitempty"`
|
|
||||||
// Model is the LLM model ID. If empty, the default model is used.
|
|
||||||
Model string `json:"model,omitempty"`
|
|
||||||
// SystemMessage overrides the configured system message for this request.
|
|
||||||
SystemMessage string `json:"system_message,omitempty"`
|
|
||||||
// MaxIterations caps the number of agent loop turns (0 = unlimited).
|
|
||||||
MaxIterations int `json:"max_iterations,omitempty"`
|
|
||||||
// ThinkingBudgetTokens is the llama.cpp reasoning budget (-1 = unlimited, 0 = off).
|
|
||||||
ThinkingBudgetTokens *int `json:"thinking_budget_tokens,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// Response is the response format for chat and voice endpoints.
|
|
||||||
type Response struct {
|
|
||||||
// Messages is the full list of messages generated during the completions request,
|
|
||||||
// including tool calls and tool results.
|
|
||||||
Messages []llm.Message `json:"messages,omitempty"`
|
|
||||||
// Provider is the LLM provider used for the response.
|
|
||||||
Provider string `json:"used_provider,omitempty"`
|
|
||||||
// Model is the LLM model used for the response.
|
|
||||||
Model string `json:"used_model,omitempty"`
|
|
||||||
// Usage is token usage from the API response.
|
|
||||||
Usage *llm.Usage `json:"usage,omitempty"`
|
|
||||||
// Timings is llama.cpp timing info from the API response.
|
|
||||||
Timings *llm.Timings `json:"timings,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// chat processes text chat requests.
|
|
||||||
func (svc *Service) chat(w http.ResponseWriter, r *http.Request) {
|
|
||||||
var (
|
|
||||||
ctx = r.Context()
|
|
||||||
log = ctx.Value(logKey).(*slog.Logger)
|
|
||||||
)
|
|
||||||
|
|
||||||
// Parse request.
|
|
||||||
r.Body = http.MaxBytesReader(w, r.Body, 32<<20)
|
|
||||||
var req Request
|
|
||||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"failed to decode request",
|
|
||||||
slog.Any("error", err),
|
|
||||||
)
|
|
||||||
http.Error(w, "invalid request", http.StatusBadRequest)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Validate messages.
|
|
||||||
if len(req.Messages) == 0 {
|
|
||||||
http.Error(w, "messages required", http.StatusBadRequest)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
log.DebugContext(
|
|
||||||
ctx,
|
|
||||||
"messages",
|
|
||||||
slog.Any("data", req.Messages),
|
|
||||||
)
|
|
||||||
|
|
||||||
// Resolve provider and model.
|
|
||||||
provider := req.Provider
|
|
||||||
if provider == "" {
|
|
||||||
provider = svc.cfg.DefaultModel.Provider
|
|
||||||
}
|
|
||||||
llmc := svc.llms[provider]
|
|
||||||
if llmc == nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"unknown provider",
|
|
||||||
slog.String("provider", provider),
|
|
||||||
)
|
|
||||||
http.Error(
|
|
||||||
w,
|
|
||||||
fmt.Sprintf("unknown provider: %q", provider),
|
|
||||||
http.StatusBadRequest,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
model := req.Model
|
|
||||||
if model == "" {
|
|
||||||
model = svc.cfg.DefaultModel.Model
|
|
||||||
}
|
|
||||||
// Validate model against provider allowed list, if configured.
|
|
||||||
if allowed, ok := svc.allowedModels[provider]; ok {
|
|
||||||
m, _, _ := strings.Cut(model, ":")
|
|
||||||
if _, ok := allowed[m]; !ok {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"model not in allowed list",
|
|
||||||
slog.String("provider", provider),
|
|
||||||
slog.String("model", model),
|
|
||||||
)
|
|
||||||
http.Error(
|
|
||||||
w,
|
|
||||||
fmt.Sprintf(
|
|
||||||
"model %q not allowed for provider %q",
|
|
||||||
model, provider,
|
|
||||||
),
|
|
||||||
http.StatusBadRequest,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
res, err := llmc.Completions(
|
|
||||||
ctx, llm.CompletionsParams{
|
|
||||||
Model: model,
|
|
||||||
SystemMessage: req.SystemMessage,
|
|
||||||
MaxIterations: req.MaxIterations,
|
|
||||||
Messages: req.Messages,
|
|
||||||
ThinkingBudgetTokens: req.ThinkingBudgetTokens,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
if err != nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"LLM request failed",
|
|
||||||
slog.Any("error", err),
|
|
||||||
)
|
|
||||||
http.Error(w, "LLM error", http.StatusInternalServerError)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if len(res.Messages) == 0 {
|
|
||||||
http.Error(
|
|
||||||
w,
|
|
||||||
"no response from LLM",
|
|
||||||
http.StatusInternalServerError,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
final := res.Messages[len(res.Messages)-1]
|
|
||||||
log.DebugContext(
|
|
||||||
ctx,
|
|
||||||
"LLM response",
|
|
||||||
slog.String("text", final.Text()),
|
|
||||||
slog.String("provider", provider),
|
|
||||||
slog.String("model", model),
|
|
||||||
)
|
|
||||||
|
|
||||||
w.Header().Set("Content-Type", "application/json")
|
|
||||||
if err := json.NewEncoder(w).Encode(Response{
|
|
||||||
Messages: res.Messages,
|
|
||||||
Provider: provider,
|
|
||||||
Model: model,
|
|
||||||
Usage: res.Usage,
|
|
||||||
Timings: res.Timings,
|
|
||||||
}); err != nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"failed to json encode response",
|
|
||||||
slog.Any("error", err),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// StreamMessage is the SSE event payload for the streaming chat endpoint.
|
|
||||||
type StreamMessage struct {
|
|
||||||
// Error is an error message, if any.
|
|
||||||
Error string `json:"error,omitempty"`
|
|
||||||
// Delta is a text fragment from the assistant response.
|
|
||||||
Delta string `json:"delta,omitempty"`
|
|
||||||
// ReasoningDelta is a reasoning text fragment from the assistant response.
|
|
||||||
ReasoningDelta string `json:"reasoning_delta,omitempty"`
|
|
||||||
// Message is the chat completion message.
|
|
||||||
Message *llm.Message `json:"message,omitempty"`
|
|
||||||
// Provider is the LLM provider used for the response.
|
|
||||||
Provider string `json:"provider,omitempty"`
|
|
||||||
// Model is the LLM model used for the response.
|
|
||||||
Model string `json:"model,omitempty"`
|
|
||||||
// Usage is token usage from the API response.
|
|
||||||
Usage *llm.Usage `json:"usage,omitempty"`
|
|
||||||
// Timings is llama.cpp timing info from the API response.
|
|
||||||
Timings *llm.Timings `json:"timings,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// chatStream processes chat requests with streaming SSE output.
|
|
||||||
func (svc *Service) chatStream(w http.ResponseWriter, r *http.Request) {
|
|
||||||
var (
|
|
||||||
ctx = r.Context()
|
|
||||||
log = ctx.Value(logKey).(*slog.Logger)
|
|
||||||
)
|
|
||||||
|
|
||||||
// Check that the response writer supports flushing.
|
|
||||||
flusher, ok := w.(http.Flusher)
|
|
||||||
if !ok {
|
|
||||||
http.Error(
|
|
||||||
w,
|
|
||||||
"streaming not supported",
|
|
||||||
http.StatusInternalServerError,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Parse request.
|
|
||||||
r.Body = http.MaxBytesReader(w, r.Body, 32<<20)
|
|
||||||
var req Request
|
|
||||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"failed to decode request",
|
|
||||||
slog.Any("error", err),
|
|
||||||
)
|
|
||||||
http.Error(w, "invalid request", http.StatusBadRequest)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Validate messages.
|
|
||||||
if len(req.Messages) == 0 {
|
|
||||||
http.Error(w, "messages required", http.StatusBadRequest)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Resolve provider and model.
|
|
||||||
provider := req.Provider
|
|
||||||
if provider == "" {
|
|
||||||
provider = svc.cfg.DefaultModel.Provider
|
|
||||||
}
|
|
||||||
model := req.Model
|
|
||||||
if model == "" {
|
|
||||||
model = svc.cfg.DefaultModel.Model
|
|
||||||
}
|
|
||||||
llmc := svc.llms[provider]
|
|
||||||
if llmc == nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"unknown provider",
|
|
||||||
slog.String("provider", provider),
|
|
||||||
)
|
|
||||||
http.Error(
|
|
||||||
w,
|
|
||||||
fmt.Sprintf("unknown provider: %q", provider),
|
|
||||||
http.StatusBadRequest,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Validate model against provider allowed list, if configured.
|
|
||||||
if allowed, ok := svc.allowedModels[provider]; ok {
|
|
||||||
m, _, _ := strings.Cut(model, ":")
|
|
||||||
if _, ok := allowed[m]; !ok {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"model not in allowed list",
|
|
||||||
slog.String("provider", provider),
|
|
||||||
slog.String("model", model),
|
|
||||||
)
|
|
||||||
http.Error(
|
|
||||||
w,
|
|
||||||
fmt.Sprintf(
|
|
||||||
"model %q not allowed for provider %q",
|
|
||||||
model, provider,
|
|
||||||
),
|
|
||||||
http.StatusBadRequest,
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Set SSE headers.
|
|
||||||
w.Header().Set("Cache-Control", "no-cache")
|
|
||||||
w.Header().Set("Connection", "keep-alive")
|
|
||||||
w.Header().Set("Content-Type", "text/event-stream")
|
|
||||||
|
|
||||||
// Helper to send an SSE event.
|
|
||||||
send := func(event string, msg StreamMessage) {
|
|
||||||
data, err := json.Marshal(msg)
|
|
||||||
if err != nil {
|
|
||||||
log.ErrorContext(ctx, "failed to marshal SSE event",
|
|
||||||
slog.Any("error", err),
|
|
||||||
)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, data)
|
|
||||||
flusher.Flush()
|
|
||||||
}
|
|
||||||
|
|
||||||
// Start streaming LLM completions request.
|
|
||||||
var (
|
|
||||||
events = make(chan llm.StreamEvent)
|
|
||||||
errs = make(chan error, 1)
|
|
||||||
)
|
|
||||||
go func() {
|
|
||||||
errs <- llmc.CompletionsStream(
|
|
||||||
ctx, llm.CompletionsParams{
|
|
||||||
Model: model,
|
|
||||||
SystemMessage: req.SystemMessage,
|
|
||||||
MaxIterations: req.MaxIterations,
|
|
||||||
Messages: req.Messages,
|
|
||||||
ThinkingBudgetTokens: req.ThinkingBudgetTokens,
|
|
||||||
}, events,
|
|
||||||
)
|
|
||||||
close(errs)
|
|
||||||
}()
|
|
||||||
|
|
||||||
// Consume events and send as SSE.
|
|
||||||
for evt := range events {
|
|
||||||
msg := StreamMessage{
|
|
||||||
Delta: evt.Delta,
|
|
||||||
ReasoningDelta: evt.ReasoningDelta,
|
|
||||||
Provider: provider,
|
|
||||||
Model: model,
|
|
||||||
}
|
|
||||||
if evt.Message.Role != "" {
|
|
||||||
msg.Message = &evt.Message
|
|
||||||
}
|
|
||||||
if evt.Usage != nil {
|
|
||||||
msg.Usage = evt.Usage
|
|
||||||
}
|
|
||||||
if evt.Timings != nil {
|
|
||||||
msg.Timings = evt.Timings
|
|
||||||
}
|
|
||||||
send(evt.Type, msg)
|
|
||||||
}
|
|
||||||
if err := <-errs; err != nil {
|
|
||||||
log.ErrorContext(
|
|
||||||
ctx,
|
|
||||||
"LLM stream failed",
|
|
||||||
slog.Any("error", err),
|
|
||||||
)
|
|
||||||
msg := llm.Message{Role: llm.RoleAssistant}
|
|
||||||
send("message", StreamMessage{
|
|
||||||
Message: &msg,
|
|
||||||
Error: fmt.Sprintf("LLM error: %v", err),
|
|
||||||
})
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// ModelStatus represents the availability status of a model.
|
// ModelStatus represents the availability status of a model.
|
||||||
type ModelStatus string
|
type ModelStatus string
|
||||||
|
|
||||||
@@ -784,14 +549,14 @@ const (
|
|||||||
ModelStatusProviderError ModelStatus = "provider_error"
|
ModelStatusProviderError ModelStatus = "provider_error"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Model represents a model in the /v1/models response.
|
// Model represents a model in the /models response.
|
||||||
type Model struct {
|
type Model struct {
|
||||||
Model string `json:"model"`
|
Model string `json:"model"`
|
||||||
Status ModelStatus `json:"status"`
|
Status ModelStatus `json:"status"`
|
||||||
ContextSize *int `json:"context_size,omitempty"`
|
ContextSize *int `json:"context_size,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Models is the response format for the /v1/models endpoint.
|
// Models is the response format for the /models endpoint.
|
||||||
type Models struct {
|
type Models struct {
|
||||||
Providers map[string][]Model `json:"providers"`
|
Providers map[string][]Model `json:"providers"`
|
||||||
DefaultProvider string `json:"default_provider"`
|
DefaultProvider string `json:"default_provider"`
|
||||||
@@ -891,7 +656,7 @@ func (svc *Service) models(w http.ResponseWriter, r *http.Request) {
|
|||||||
wg.Wait()
|
wg.Wait()
|
||||||
|
|
||||||
w.Header().Set("Content-Type", "application/json")
|
w.Header().Set("Content-Type", "application/json")
|
||||||
if err := json.NewEncoder(w).Encode(Models{
|
if err := json.MarshalWrite(w, Models{
|
||||||
Providers: providers,
|
Providers: providers,
|
||||||
DefaultProvider: svc.cfg.DefaultModel.Provider,
|
DefaultProvider: svc.cfg.DefaultModel.Provider,
|
||||||
DefaultModel: svc.cfg.DefaultModel.Model,
|
DefaultModel: svc.cfg.DefaultModel.Model,
|
||||||
|
|||||||
Reference in New Issue
Block a user