3963eede70
Go reverse proxy for ollama.com that balances requests across multiple API keys with round-robin and failover on 429/5xx. Exposes both native Ollama API (/api/*) and OpenAI-compatible (/v1/*) passthrough. - Round-robin balancer with per-account cooldown (60s default) - Pre-stream failover: 429 → cooldown + next account; 5xx → next account - Streaming invariant: once 2xx starts streaming, no account switch - SSE (text/event-stream) and NDJSON passthrough with http.Flusher - CLI: accounts add/list/remove/set-base-url, serve, version - Accounts stored in ~/.config/ollama-proxy/accounts.json (chmod 0600) - systemd unit (User=dueattendant149, 127.0.0.1:11435, Restart=always) - 43 tests (unit + integration with httptest upstream) - opencode integration: custom provider 'ocp' with explicit model list - Requires NO_PROXY=127.0.0.1,localhost when HTTP_PROXY is set
109 lines
2.9 KiB
Go
109 lines
2.9 KiB
Go
package cli
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/Atte149/ollama-proxy/internal/config"
|
|
"github.com/Atte149/ollama-proxy/internal/log"
|
|
"github.com/Atte149/ollama-proxy/internal/proxy"
|
|
)
|
|
|
|
// RunServe starts the proxy HTTP server.
|
|
func RunServe(args []string) int {
|
|
fs := flag.NewFlagSet("serve", flag.ContinueOnError)
|
|
addr := fs.String("addr", "", "listen address (default 127.0.0.1:11435)")
|
|
baseURL := fs.String("base-url", "", "upstream root (default https://ollama.com)")
|
|
cooldown := fs.Duration("cooldown", 0, "per-account 429 cooldown window")
|
|
retries := fs.Int("retries", 0, "max failover attempts")
|
|
logLevel := fs.String("log-level", "", "debug|info|warn|error")
|
|
if err := fs.Parse(args); err != nil {
|
|
fmt.Fprintln(os.Stderr, err)
|
|
return 2
|
|
}
|
|
|
|
cfg := config.DefaultServerConfig()
|
|
if *addr != "" {
|
|
cfg.Addr = *addr
|
|
}
|
|
if *baseURL != "" {
|
|
cfg.BaseURL = *baseURL
|
|
}
|
|
if *cooldown > 0 {
|
|
cfg.Cooldown = *cooldown
|
|
}
|
|
if *retries > 0 {
|
|
cfg.Retries = *retries
|
|
}
|
|
if *logLevel != "" {
|
|
cfg.LogLevel = *logLevel
|
|
}
|
|
cfg.ApplyEnv()
|
|
if err := cfg.Validate(); err != nil {
|
|
fmt.Fprintln(os.Stderr, "config error:", err)
|
|
return 2
|
|
}
|
|
|
|
logger := log.New(cfg.LogLevel)
|
|
|
|
af, err := config.LoadAccounts()
|
|
if err != nil {
|
|
logger.Error("load accounts failed", "err", err)
|
|
return 1
|
|
}
|
|
if len(af.Accounts) == 0 {
|
|
logger.Error("no accounts configured; run `ollama-proxy accounts add` first")
|
|
return 1
|
|
}
|
|
// Per-file BaseURL override wins over flag/env when set.
|
|
if af.BaseURL != "" {
|
|
cfg.BaseURL = af.BaseURL
|
|
}
|
|
|
|
balancer := proxy.NewBalancer(af.Accounts, cfg.Cooldown)
|
|
handler := proxy.NewHandler(balancer, cfg.BaseURL, cfg.Retries, logger)
|
|
|
|
srv := &http.Server{
|
|
Addr: cfg.Addr,
|
|
Handler: handler,
|
|
// No ReadTimeout/WriteTimeout: streaming chat may be idle for long
|
|
// periods between tokens. The http.Server's default idle timeout (60s)
|
|
// applies between reads; if that becomes a problem we'll add a per-read
|
|
// deadline via a wrapped Transport, not a global timeout.
|
|
}
|
|
|
|
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer stop()
|
|
|
|
go func() {
|
|
logger.Info("ollama-proxy starting",
|
|
"addr", cfg.Addr, "upstream", cfg.BaseURL,
|
|
"accounts", balancer.Len(), "cooldown", cfg.Cooldown.String(),
|
|
"retries", cfg.Retries, "log_level", cfg.LogLevel)
|
|
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
logger.Error("server error", "err", err)
|
|
stop()
|
|
}
|
|
}()
|
|
|
|
<-ctx.Done()
|
|
logger.Info("shutdown signal received, draining…")
|
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
if err := srv.Shutdown(shutdownCtx); err != nil {
|
|
logger.Error("graceful shutdown failed", "err", err)
|
|
return 1
|
|
}
|
|
logger.Info("stopped")
|
|
return 0
|
|
}
|
|
|
|
// keep slog referenced even if logger.go ever drops it
|
|
var _ = slog.LevelInfo
|