Fix service shutdown
This commit is contained in:
+148
-383
@@ -6,7 +6,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"encoding/json"
|
||||
"encoding/json/v2"
|
||||
"fmt"
|
||||
"html/template"
|
||||
"io/fs"
|
||||
@@ -19,13 +19,16 @@ import (
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
"uuid"
|
||||
|
||||
"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/llm"
|
||||
"code.chimeric.al/chimerical/odidere/internal/service/templates"
|
||||
"code.chimeric.al/chimerical/odidere/internal/tool"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/robfig/cron/v3"
|
||||
)
|
||||
|
||||
@@ -48,6 +51,9 @@ var static embed.FS
|
||||
type Service struct {
|
||||
cfg *config.Config
|
||||
cron *cron.Cron
|
||||
db *db.SQLite
|
||||
dbr *db.Queries
|
||||
dbw *db.Queries
|
||||
jobs []*job.Job
|
||||
llms map[string]*llm.Client
|
||||
log *slog.Logger
|
||||
@@ -55,20 +61,49 @@ type Service struct {
|
||||
mux *http.ServeMux
|
||||
server *http.Server
|
||||
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.
|
||||
// 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{
|
||||
cfg: cfg,
|
||||
db: sdb,
|
||||
dbr: db.New(sdb.Reader),
|
||||
dbw: db.New(sdb.Writer),
|
||||
log: log,
|
||||
mux: http.NewServeMux(),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
conversations: make(map[string]*conversation.Conversation),
|
||||
}
|
||||
|
||||
// Setup tool registry.
|
||||
registry, err := llm.NewRegistry(cfg.Tools)
|
||||
registry, err := tool.NewRegistry(cfg.Tools)
|
||||
if err != nil {
|
||||
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(
|
||||
llm.Config{
|
||||
Key: pc.Key,
|
||||
SystemMessage: cfg.SystemMessage,
|
||||
Provider: pc.Name,
|
||||
Timeout: pc.Timeout,
|
||||
URL: pc.URL,
|
||||
},
|
||||
registry,
|
||||
log,
|
||||
)
|
||||
if err != nil {
|
||||
@@ -100,6 +134,9 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
||||
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.
|
||||
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))
|
||||
for _, jc := range cfg.Jobs {
|
||||
// Resolve the provider and model for this job.
|
||||
var provider, modelID string
|
||||
if jc.Model.Provider != "" {
|
||||
if jc.Model.Model == "" {
|
||||
return nil, fmt.Errorf(
|
||||
"job %q: model.provider set but model.model is empty",
|
||||
jc.Name,
|
||||
)
|
||||
provider, modelID := jc.Model.Provider, jc.Model.Model
|
||||
if provider == "" {
|
||||
provider = cfg.DefaultModel.Provider
|
||||
}
|
||||
provider = jc.Model.Provider
|
||||
modelID = jc.Model.Model
|
||||
} else {
|
||||
provider = svc.cfg.DefaultModel.Provider
|
||||
modelID = svc.cfg.DefaultModel.Model
|
||||
if modelID == "" {
|
||||
modelID = cfg.DefaultModel.Model
|
||||
}
|
||||
llmc := svc.llms[provider]
|
||||
if llmc == nil {
|
||||
@@ -143,7 +173,14 @@ func New(cfg *config.Config, log *slog.Logger) (*Service, error) {
|
||||
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 {
|
||||
return nil, fmt.Errorf(
|
||||
"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)),
|
||||
),
|
||||
)
|
||||
svc.mux.HandleFunc("POST /v1/chat/voice", svc.chat)
|
||||
svc.mux.HandleFunc("POST /v1/chat/voice/stream", svc.chatStream)
|
||||
svc.mux.HandleFunc("GET /v1/models", svc.models)
|
||||
svc.mux.HandleFunc("GET /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{
|
||||
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) {
|
||||
var (
|
||||
start = time.Now()
|
||||
id = uuid.NewString()
|
||||
id = uuid.NewV4().String()
|
||||
ip = func() string {
|
||||
if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
|
||||
before, _, found := strings.Cut(ip, ",")
|
||||
@@ -296,7 +370,7 @@ func (svc *Service) Run(ctx context.Context) error {
|
||||
"jobs",
|
||||
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.
|
||||
@@ -305,26 +379,29 @@ func (svc *Service) Run(ctx context.Context) error {
|
||||
os.Interrupt, syscall.SIGTERM,
|
||||
)
|
||||
defer cancel()
|
||||
svc.ctx, svc.cancel = context.WithCancel(ctx)
|
||||
defer svc.cancel()
|
||||
|
||||
// Register jobs with cron scheduler.
|
||||
svc.cron = cron.New()
|
||||
for _, j := range svc.jobs {
|
||||
entry, err := svc.cron.AddFunc(j.Schedule(), func() {
|
||||
result := j.Run(ctx)
|
||||
if result.Error != nil {
|
||||
title := j.Name()
|
||||
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",
|
||||
slog.String("name", j.Name()),
|
||||
slog.Any("error", result.Error),
|
||||
slog.Duration(
|
||||
"duration", result.Duration,
|
||||
),
|
||||
)
|
||||
} else {
|
||||
svc.log.Info("job completed",
|
||||
slog.String("name", j.Name()),
|
||||
slog.Duration(
|
||||
"duration", result.Duration,
|
||||
),
|
||||
slog.Any("error", err),
|
||||
)
|
||||
}
|
||||
})
|
||||
@@ -358,36 +435,56 @@ func (svc *Service) Run(ctx context.Context) error {
|
||||
}()
|
||||
|
||||
// Wait for shutdown signal or server error.
|
||||
var runErr error
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
svc.log.Info("shutdown signal received")
|
||||
case err := <-errs:
|
||||
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(
|
||||
context.Background(),
|
||||
svc.cfg.GetShutdownTimeout(),
|
||||
svc.cfg.ShutdownTimeout,
|
||||
)
|
||||
defer shutdownCancel()
|
||||
|
||||
// Stop cron scheduler — in-flight jobs continue to completion.
|
||||
svc.log.Info("stopping job scheduler")
|
||||
svc.cron.Stop()
|
||||
svc.cancel()
|
||||
|
||||
svc.log.Info("stopping job scheduler")
|
||||
cronCtx := svc.cron.Stop()
|
||||
|
||||
var shutdownErr error
|
||||
svc.log.Info("shutting down HTTP server")
|
||||
if err := svc.server.Shutdown(shutdownCtx); err != nil {
|
||||
svc.log.Warn(
|
||||
"shutdown timeout reached",
|
||||
slog.Any("error", err),
|
||||
shutdownErr = fmt.Errorf("HTTP shutdown: %w", err)
|
||||
} else {
|
||||
// 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")
|
||||
return nil
|
||||
return runErr
|
||||
}
|
||||
|
||||
// ListJobs returns all registered jobs.
|
||||
@@ -438,338 +535,6 @@ func (svc *Service) status(w http.ResponseWriter, r *http.Request) {
|
||||
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.
|
||||
type ModelStatus string
|
||||
|
||||
@@ -784,14 +549,14 @@ const (
|
||||
ModelStatusProviderError ModelStatus = "provider_error"
|
||||
)
|
||||
|
||||
// Model represents a model in the /v1/models response.
|
||||
// Model represents a model in the /models response.
|
||||
type Model struct {
|
||||
Model string `json:"model"`
|
||||
Status ModelStatus `json:"status"`
|
||||
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 {
|
||||
Providers map[string][]Model `json:"providers"`
|
||||
DefaultProvider string `json:"default_provider"`
|
||||
@@ -891,7 +656,7 @@ func (svc *Service) models(w http.ResponseWriter, r *http.Request) {
|
||||
wg.Wait()
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if err := json.NewEncoder(w).Encode(Models{
|
||||
if err := json.MarshalWrite(w, Models{
|
||||
Providers: providers,
|
||||
DefaultProvider: svc.cfg.DefaultModel.Provider,
|
||||
DefaultModel: svc.cfg.DefaultModel.Model,
|
||||
|
||||
Reference in New Issue
Block a user