+
+
package main
+
+import (
+ "context"
+ "database/sql"
+ "flag"
+ "fmt"
+ "log"
+ "log/slog"
+ "net/http"
+ "os"
+ "os/signal"
+ "syscall"
+ "time"
+
+ _ "github.com/go-sql-driver/mysql"
+ "github.com/google/uuid"
+ _ "github.com/jackc/pgx/v5/stdlib"
+ _ "github.com/mattn/go-sqlite3"
+ "github.com/redis/go-redis/v9"
+
+ "github.com/ajac-zero/latticelm/internal/auth"
+ "github.com/ajac-zero/latticelm/internal/config"
+ "github.com/ajac-zero/latticelm/internal/conversation"
+ slogger "github.com/ajac-zero/latticelm/internal/logger"
+ "github.com/ajac-zero/latticelm/internal/observability"
+ "github.com/ajac-zero/latticelm/internal/providers"
+ "github.com/ajac-zero/latticelm/internal/ratelimit"
+ "github.com/ajac-zero/latticelm/internal/server"
+ "github.com/prometheus/client_golang/prometheus"
+ "github.com/prometheus/client_golang/prometheus/promhttp"
+ "go.opentelemetry.io/otel"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+)
+
+func main() {
+ var configPath string
+ flag.StringVar(&configPath, "config", "config.yaml", "path to config file")
+ flag.Parse()
+
+ cfg, err := config.Load(configPath)
+ if err != nil {
+ log.Fatalf("load config: %v", err)
+ }
+
+ // Initialize logger from config
+ logFormat := cfg.Logging.Format
+ if logFormat == "" {
+ logFormat = "json"
+ }
+ logLevel := cfg.Logging.Level
+ if logLevel == "" {
+ logLevel = "info"
+ }
+ logger := slogger.New(logFormat, logLevel)
+
+ // Initialize tracing
+ var tracerProvider *sdktrace.TracerProvider
+ if cfg.Observability.Enabled && cfg.Observability.Tracing.Enabled {
+ // Set defaults
+ tracingCfg := cfg.Observability.Tracing
+ if tracingCfg.ServiceName == "" {
+ tracingCfg.ServiceName = "llm-gateway"
+ }
+ if tracingCfg.Sampler.Type == "" {
+ tracingCfg.Sampler.Type = "probability"
+ tracingCfg.Sampler.Rate = 0.1
+ }
+
+ tp, err := observability.InitTracer(tracingCfg)
+ if err != nil {
+ logger.Error("failed to initialize tracing", slog.String("error", err.Error()))
+ } else {
+ tracerProvider = tp
+ otel.SetTracerProvider(tracerProvider)
+ logger.Info("tracing initialized",
+ slog.String("exporter", tracingCfg.Exporter.Type),
+ slog.String("sampler", tracingCfg.Sampler.Type),
+ )
+ }
+ }
+
+ // Initialize metrics
+ var metricsRegistry *prometheus.Registry
+ if cfg.Observability.Enabled && cfg.Observability.Metrics.Enabled {
+ metricsRegistry = observability.InitMetrics()
+ metricsPath := cfg.Observability.Metrics.Path
+ if metricsPath == "" {
+ metricsPath = "/metrics"
+ }
+ logger.Info("metrics initialized", slog.String("path", metricsPath))
+ }
+
+ // Create provider registry with circuit breaker support
+ var baseRegistry *providers.Registry
+ if cfg.Observability.Enabled && cfg.Observability.Metrics.Enabled {
+ // Pass observability callback for circuit breaker state changes
+ baseRegistry, err = providers.NewRegistryWithCircuitBreaker(
+ cfg.Providers,
+ cfg.Models,
+ observability.RecordCircuitBreakerStateChange,
+ )
+ } else {
+ // No observability, use default registry
+ baseRegistry, err = providers.NewRegistry(cfg.Providers, cfg.Models)
+ }
+ if err != nil {
+ logger.Error("failed to initialize providers", slog.String("error", err.Error()))
+ os.Exit(1)
+ }
+
+ // Wrap providers with observability
+ var registry server.ProviderRegistry = baseRegistry
+ if cfg.Observability.Enabled {
+ registry = observability.WrapProviderRegistry(registry, metricsRegistry, tracerProvider)
+ logger.Info("providers instrumented")
+ }
+
+ // Initialize authentication middleware
+ authConfig := auth.Config{
+ Enabled: cfg.Auth.Enabled,
+ Issuer: cfg.Auth.Issuer,
+ Audience: cfg.Auth.Audience,
+ }
+ authMiddleware, err := auth.New(authConfig, logger)
+ if err != nil {
+ logger.Error("failed to initialize auth", slog.String("error", err.Error()))
+ os.Exit(1)
+ }
+
+ if cfg.Auth.Enabled {
+ logger.Info("authentication enabled", slog.String("issuer", cfg.Auth.Issuer))
+ } else {
+ logger.Warn("authentication disabled - API is publicly accessible")
+ }
+
+ // Initialize conversation store
+ convStore, storeBackend, err := initConversationStore(cfg.Conversations, logger)
+ if err != nil {
+ logger.Error("failed to initialize conversation store", slog.String("error", err.Error()))
+ os.Exit(1)
+ }
+
+ // Wrap conversation store with observability
+ if cfg.Observability.Enabled && convStore != nil {
+ convStore = observability.WrapConversationStore(convStore, storeBackend, metricsRegistry, tracerProvider)
+ logger.Info("conversation store instrumented")
+ }
+
+ gatewayServer := server.New(registry, convStore, logger)
+ mux := http.NewServeMux()
+ gatewayServer.RegisterRoutes(mux)
+
+ // Register metrics endpoint if enabled
+ if cfg.Observability.Enabled && cfg.Observability.Metrics.Enabled {
+ metricsPath := cfg.Observability.Metrics.Path
+ if metricsPath == "" {
+ metricsPath = "/metrics"
+ }
+ mux.Handle(metricsPath, promhttp.HandlerFor(metricsRegistry, promhttp.HandlerOpts{}))
+ logger.Info("metrics endpoint registered", slog.String("path", metricsPath))
+ }
+
+ addr := cfg.Server.Address
+ if addr == "" {
+ addr = ":8080"
+ }
+
+ // Initialize rate limiting
+ rateLimitConfig := ratelimit.Config{
+ Enabled: cfg.RateLimit.Enabled,
+ RequestsPerSecond: cfg.RateLimit.RequestsPerSecond,
+ Burst: cfg.RateLimit.Burst,
+ }
+ // Set defaults if not configured
+ if rateLimitConfig.Enabled && rateLimitConfig.RequestsPerSecond == 0 {
+ rateLimitConfig.RequestsPerSecond = 10 // default 10 req/s
+ }
+ if rateLimitConfig.Enabled && rateLimitConfig.Burst == 0 {
+ rateLimitConfig.Burst = 20 // default burst of 20
+ }
+ rateLimitMiddleware := ratelimit.New(rateLimitConfig, logger)
+
+ if cfg.RateLimit.Enabled {
+ logger.Info("rate limiting enabled",
+ slog.Float64("requests_per_second", rateLimitConfig.RequestsPerSecond),
+ slog.Int("burst", rateLimitConfig.Burst),
+ )
+ }
+
+ // Determine max request body size
+ maxRequestBodySize := cfg.Server.MaxRequestBodySize
+ if maxRequestBodySize == 0 {
+ maxRequestBodySize = server.MaxRequestBodyBytes // default: 10MB
+ }
+
+ logger.Info("server configuration",
+ slog.Int64("max_request_body_bytes", maxRequestBodySize),
+ )
+
+ // Build handler chain: panic recovery -> request size limit -> logging -> tracing -> metrics -> rate limiting -> auth -> routes
+ handler := server.PanicRecoveryMiddleware(
+ server.RequestSizeLimitMiddleware(
+ loggingMiddleware(
+ observability.TracingMiddleware(
+ observability.MetricsMiddleware(
+ rateLimitMiddleware.Handler(authMiddleware.Handler(mux)),
+ metricsRegistry,
+ tracerProvider,
+ ),
+ tracerProvider,
+ ),
+ logger,
+ ),
+ maxRequestBodySize,
+ ),
+ logger,
+ )
+
+ srv := &http.Server{
+ Addr: addr,
+ Handler: handler,
+ ReadTimeout: 15 * time.Second,
+ WriteTimeout: 60 * time.Second,
+ IdleTimeout: 120 * time.Second,
+ }
+
+ // Set up signal handling for graceful shutdown
+ sigChan := make(chan os.Signal, 1)
+ signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+
+ // Run server in a goroutine
+ serverErrors := make(chan error, 1)
+ go func() {
+ logger.Info("open responses gateway listening", slog.String("address", addr))
+ serverErrors <- srv.ListenAndServe()
+ }()
+
+ // Wait for shutdown signal or server error
+ select {
+ case err := <-serverErrors:
+ if err != nil && err != http.ErrServerClosed {
+ logger.Error("server error", slog.String("error", err.Error()))
+ os.Exit(1)
+ }
+ case sig := <-sigChan:
+ logger.Info("received shutdown signal", slog.String("signal", sig.String()))
+
+ // Create shutdown context with timeout
+ shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
+ defer shutdownCancel()
+
+ // Shutdown the HTTP server gracefully
+ logger.Info("shutting down server gracefully")
+ if err := srv.Shutdown(shutdownCtx); err != nil {
+ logger.Error("server shutdown error", slog.String("error", err.Error()))
+ }
+
+ // Shutdown tracer provider
+ if tracerProvider != nil {
+ logger.Info("shutting down tracer")
+ shutdownTracerCtx, shutdownTracerCancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer shutdownTracerCancel()
+ if err := observability.Shutdown(shutdownTracerCtx, tracerProvider); err != nil {
+ logger.Error("error shutting down tracer", slog.String("error", err.Error()))
+ }
+ }
+
+ // Close conversation store
+ logger.Info("closing conversation store")
+ if err := convStore.Close(); err != nil {
+ logger.Error("error closing conversation store", slog.String("error", err.Error()))
+ }
+
+ logger.Info("shutdown complete")
+ }
+}
+
+func initConversationStore(cfg config.ConversationConfig, logger *slog.Logger) (conversation.Store, string, error) {
+ var ttl time.Duration
+ if cfg.TTL != "" {
+ parsed, err := time.ParseDuration(cfg.TTL)
+ if err != nil {
+ return nil, "", fmt.Errorf("invalid conversation ttl %q: %w", cfg.TTL, err)
+ }
+ ttl = parsed
+ }
+
+ switch cfg.Store {
+ case "sql":
+ driver := cfg.Driver
+ if driver == "" {
+ driver = "sqlite3"
+ }
+ db, err := sql.Open(driver, cfg.DSN)
+ if err != nil {
+ return nil, "", fmt.Errorf("open database: %w", err)
+ }
+ store, err := conversation.NewSQLStore(db, driver, ttl)
+ if err != nil {
+ return nil, "", fmt.Errorf("init sql store: %w", err)
+ }
+ logger.Info("conversation store initialized",
+ slog.String("backend", "sql"),
+ slog.String("driver", driver),
+ slog.Duration("ttl", ttl),
+ )
+ return store, "sql", nil
+ case "redis":
+ opts, err := redis.ParseURL(cfg.DSN)
+ if err != nil {
+ return nil, "", fmt.Errorf("parse redis dsn: %w", err)
+ }
+ client := redis.NewClient(opts)
+
+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer cancel()
+
+ if err := client.Ping(ctx).Err(); err != nil {
+ return nil, "", fmt.Errorf("connect to redis: %w", err)
+ }
+
+ logger.Info("conversation store initialized",
+ slog.String("backend", "redis"),
+ slog.Duration("ttl", ttl),
+ )
+ return conversation.NewRedisStore(client, ttl), "redis", nil
+ default:
+ logger.Info("conversation store initialized",
+ slog.String("backend", "memory"),
+ slog.Duration("ttl", ttl),
+ )
+ return conversation.NewMemoryStore(ttl), "memory", nil
+ }
+}
+type responseWriter struct {
+ http.ResponseWriter
+ statusCode int
+ bytesWritten int
+}
+
+func (rw *responseWriter) WriteHeader(code int) {
+ rw.statusCode = code
+ rw.ResponseWriter.WriteHeader(code)
+}
+
+func (rw *responseWriter) Write(b []byte) (int, error) {
+ n, err := rw.ResponseWriter.Write(b)
+ rw.bytesWritten += n
+ return n, err
+}
+
+func loggingMiddleware(next http.Handler, logger *slog.Logger) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ start := time.Now()
+
+ // Generate request ID
+ requestID := uuid.NewString()
+ ctx := slogger.WithRequestID(r.Context(), requestID)
+ r = r.WithContext(ctx)
+
+ // Wrap response writer to capture status code
+ rw := &responseWriter{
+ ResponseWriter: w,
+ statusCode: http.StatusOK,
+ }
+
+ // Add request ID header
+ w.Header().Set("X-Request-ID", requestID)
+
+ // Log request start
+ logger.InfoContext(ctx, "request started",
+ slog.String("request_id", requestID),
+ slog.String("method", r.Method),
+ slog.String("path", r.URL.Path),
+ slog.String("remote_addr", r.RemoteAddr),
+ slog.String("user_agent", r.UserAgent()),
+ )
+
+ next.ServeHTTP(rw, r)
+
+ duration := time.Since(start)
+
+ // Log request completion with appropriate level
+ logLevel := slog.LevelInfo
+ if rw.statusCode >= 500 {
+ logLevel = slog.LevelError
+ } else if rw.statusCode >= 400 {
+ logLevel = slog.LevelWarn
+ }
+
+ logger.Log(ctx, logLevel, "request completed",
+ slog.String("request_id", requestID),
+ slog.String("method", r.Method),
+ slog.String("path", r.URL.Path),
+ slog.Int("status_code", rw.statusCode),
+ slog.Int("response_bytes", rw.bytesWritten),
+ slog.Duration("duration", duration),
+ slog.Float64("duration_ms", float64(duration.Milliseconds())),
+ )
+ })
+}
+
+
+
package api
+
+import (
+ "encoding/json"
+ "errors"
+ "fmt"
+)
+
+// ============================================================
+// Request Types (CreateResponseBody)
+// ============================================================
+
+// ResponseRequest models the OpenResponses CreateResponseBody.
+type ResponseRequest struct {
+ Model string `json:"model"`
+ Input InputUnion `json:"input"`
+ Instructions *string `json:"instructions,omitempty"`
+ MaxOutputTokens *int `json:"max_output_tokens,omitempty"`
+ Metadata map[string]string `json:"metadata,omitempty"`
+ Stream bool `json:"stream,omitempty"`
+ PreviousResponseID *string `json:"previous_response_id,omitempty"`
+ Temperature *float64 `json:"temperature,omitempty"`
+ TopP *float64 `json:"top_p,omitempty"`
+ FrequencyPenalty *float64 `json:"frequency_penalty,omitempty"`
+ PresencePenalty *float64 `json:"presence_penalty,omitempty"`
+ TopLogprobs *int `json:"top_logprobs,omitempty"`
+ Truncation *string `json:"truncation,omitempty"`
+ ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
+ Tools json.RawMessage `json:"tools,omitempty"`
+ ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
+ Store *bool `json:"store,omitempty"`
+ Text json.RawMessage `json:"text,omitempty"`
+ Reasoning json.RawMessage `json:"reasoning,omitempty"`
+ Include []string `json:"include,omitempty"`
+ ServiceTier *string `json:"service_tier,omitempty"`
+ Background *bool `json:"background,omitempty"`
+ StreamOptions json.RawMessage `json:"stream_options,omitempty"`
+ MaxToolCalls *int `json:"max_tool_calls,omitempty"`
+
+ // Non-spec extension: allows client to select a specific provider.
+ Provider string `json:"provider,omitempty"`
+}
+
+// InputUnion handles the polymorphic "input" field: string or []InputItem.
+type InputUnion struct {
+ String *string
+ Items []InputItem
+}
+
+func (u *InputUnion) UnmarshalJSON(data []byte) error {
+ if string(data) == "null" {
+ return nil
+ }
+ var s string
+ if err := json.Unmarshal(data, &s); err == nil {
+ u.String = &s
+ return nil
+ }
+ var items []InputItem
+ if err := json.Unmarshal(data, &items); err == nil {
+ u.Items = items
+ return nil
+ }
+ return fmt.Errorf("input must be a string or array of items")
+}
+
+func (u InputUnion) MarshalJSON() ([]byte, error) {
+ if u.String != nil {
+ return json.Marshal(*u.String)
+ }
+ if u.Items != nil {
+ return json.Marshal(u.Items)
+ }
+ return []byte("null"), nil
+}
+
+// InputItem is a discriminated union on "type".
+// Valid types: message, item_reference, function_call, function_call_output, reasoning.
+type InputItem struct {
+ Type string `json:"type"`
+ Role string `json:"role,omitempty"`
+ Content json.RawMessage `json:"content,omitempty"`
+ ID string `json:"id,omitempty"`
+ CallID string `json:"call_id,omitempty"`
+ Name string `json:"name,omitempty"`
+ Arguments string `json:"arguments,omitempty"`
+ Output string `json:"output,omitempty"`
+ Status string `json:"status,omitempty"`
+}
+
+// ============================================================
+// Internal Types (providers + conversation store)
+// ============================================================
+
+// Message is the normalized internal message representation.
+type Message struct {
+ Role string `json:"role"`
+ Content []ContentBlock `json:"content"`
+ CallID string `json:"call_id,omitempty"` // for tool messages
+ Name string `json:"name,omitempty"` // for tool messages
+ ToolCalls []ToolCall `json:"tool_calls,omitempty"` // for assistant messages
+}
+
+// ContentBlock is a typed content element.
+type ContentBlock struct {
+ Type string `json:"type"`
+ Text string `json:"text,omitempty"`
+}
+
+// NormalizeInput converts the request Input into messages for providers.
+// Does NOT include instructions (the server prepends those separately).
+func (r *ResponseRequest) NormalizeInput() []Message {
+ if r.Input.String != nil {
+ return []Message{{
+ Role: "user",
+ Content: []ContentBlock{{Type: "input_text", Text: *r.Input.String}},
+ }}
+ }
+
+ var msgs []Message
+ for _, item := range r.Input.Items {
+ switch item.Type {
+ case "message", "":
+ msg := Message{Role: item.Role}
+ if item.Content != nil {
+ var s string
+ if err := json.Unmarshal(item.Content, &s); err == nil {
+ contentType := "input_text"
+ if item.Role == "assistant" {
+ contentType = "output_text"
+ }
+ msg.Content = []ContentBlock{{Type: contentType, Text: s}}
+ } else {
+ // Content is an array of blocks - parse them
+ var rawBlocks []map[string]interface{}
+ if err := json.Unmarshal(item.Content, &rawBlocks); err == nil {
+ // Extract content blocks and tool calls
+ for _, block := range rawBlocks {
+ blockType, _ := block["type"].(string)
+
+ if blockType == "tool_use" {
+ // Extract tool call information
+ toolCall := ToolCall{
+ ID: getStringField(block, "id"),
+ Name: getStringField(block, "name"),
+ }
+ // input field contains the arguments as a map
+ if input, ok := block["input"].(map[string]interface{}); ok {
+ if inputJSON, err := json.Marshal(input); err == nil {
+ toolCall.Arguments = string(inputJSON)
+ }
+ }
+ msg.ToolCalls = append(msg.ToolCalls, toolCall)
+ } else if blockType == "output_text" || blockType == "input_text" {
+ // Regular text content block
+ msg.Content = append(msg.Content, ContentBlock{
+ Type: blockType,
+ Text: getStringField(block, "text"),
+ })
+ }
+ }
+ }
+ }
+ }
+ msgs = append(msgs, msg)
+ case "function_call_output":
+ msgs = append(msgs, Message{
+ Role: "tool",
+ Content: []ContentBlock{{Type: "input_text", Text: item.Output}},
+ CallID: item.CallID,
+ Name: item.Name,
+ })
+ }
+ }
+ return msgs
+}
+
+// ============================================================
+// Response Types (ResponseResource)
+// ============================================================
+
+// Response is the spec-compliant ResponseResource.
+type Response struct {
+ ID string `json:"id"`
+ Object string `json:"object"`
+ CreatedAt int64 `json:"created_at"`
+ CompletedAt *int64 `json:"completed_at"`
+ Status string `json:"status"`
+ IncompleteDetails *IncompleteDetails `json:"incomplete_details"`
+ Model string `json:"model"`
+ PreviousResponseID *string `json:"previous_response_id"`
+ Instructions *string `json:"instructions"`
+ Output []OutputItem `json:"output"`
+ Error *ResponseError `json:"error"`
+ Tools json.RawMessage `json:"tools"`
+ ToolChoice json.RawMessage `json:"tool_choice"`
+ Truncation string `json:"truncation"`
+ ParallelToolCalls bool `json:"parallel_tool_calls"`
+ Text json.RawMessage `json:"text"`
+ TopP float64 `json:"top_p"`
+ PresencePenalty float64 `json:"presence_penalty"`
+ FrequencyPenalty float64 `json:"frequency_penalty"`
+ TopLogprobs int `json:"top_logprobs"`
+ Temperature float64 `json:"temperature"`
+ Reasoning json.RawMessage `json:"reasoning"`
+ Usage *Usage `json:"usage"`
+ MaxOutputTokens *int `json:"max_output_tokens"`
+ MaxToolCalls *int `json:"max_tool_calls"`
+ Store bool `json:"store"`
+ Background bool `json:"background"`
+ ServiceTier string `json:"service_tier"`
+ Metadata map[string]string `json:"metadata"`
+ SafetyIdentifier *string `json:"safety_identifier"`
+ PromptCacheKey *string `json:"prompt_cache_key"`
+
+ // Non-spec extension
+ Provider string `json:"provider,omitempty"`
+}
+
+// OutputItem represents a typed item in the response output.
+type OutputItem struct {
+ ID string `json:"id"`
+ Type string `json:"type"`
+ Status string `json:"status"`
+ Role string `json:"role,omitempty"`
+ Content []ContentPart `json:"content,omitempty"`
+ CallID string `json:"call_id,omitempty"` // for function_call
+ Name string `json:"name,omitempty"` // for function_call
+ Arguments string `json:"arguments,omitempty"` // for function_call
+}
+
+// ContentPart is a content block within an output item.
+type ContentPart struct {
+ Type string `json:"type"`
+ Text string `json:"text"`
+ Annotations []Annotation `json:"annotations"`
+}
+
+// Annotation on output text content.
+type Annotation struct {
+ Type string `json:"type"`
+}
+
+// IncompleteDetails explains why a response is incomplete.
+type IncompleteDetails struct {
+ Reason string `json:"reason"`
+}
+
+// ResponseError describes an error in the response.
+type ResponseError struct {
+ Type string `json:"type"`
+ Message string `json:"message"`
+ Code *string `json:"code"`
+}
+
+// ============================================================
+// Usage Types
+// ============================================================
+
+// Usage captures token accounting with sub-details.
+type Usage struct {
+ InputTokens int `json:"input_tokens"`
+ OutputTokens int `json:"output_tokens"`
+ TotalTokens int `json:"total_tokens"`
+ InputTokensDetails InputTokensDetails `json:"input_tokens_details"`
+ OutputTokensDetails OutputTokensDetails `json:"output_tokens_details"`
+}
+
+// InputTokensDetails breaks down input token usage.
+type InputTokensDetails struct {
+ CachedTokens int `json:"cached_tokens"`
+}
+
+// OutputTokensDetails breaks down output token usage.
+type OutputTokensDetails struct {
+ ReasoningTokens int `json:"reasoning_tokens"`
+}
+
+// ============================================================
+// Streaming Types
+// ============================================================
+
+// StreamEvent represents a single SSE event in the streaming response.
+// Fields are selectively populated based on the event Type.
+type StreamEvent struct {
+ Type string `json:"type"`
+ SequenceNumber int `json:"sequence_number"`
+ Response *Response `json:"response,omitempty"`
+ OutputIndex *int `json:"output_index,omitempty"`
+ Item *OutputItem `json:"item,omitempty"`
+ ItemID string `json:"item_id,omitempty"`
+ ContentIndex *int `json:"content_index,omitempty"`
+ Part *ContentPart `json:"part,omitempty"`
+ Delta string `json:"delta,omitempty"`
+ Text string `json:"text,omitempty"`
+ Arguments string `json:"arguments,omitempty"` // for function_call_arguments.done
+}
+
+// ============================================================
+// Provider Result Types (internal, not exposed via HTTP)
+// ============================================================
+
+// ProviderResult is returned by Provider.Generate.
+type ProviderResult struct {
+ ID string
+ Model string
+ Text string
+ Usage Usage
+ ToolCalls []ToolCall
+}
+
+// ProviderStreamDelta is sent through the stream channel.
+type ProviderStreamDelta struct {
+ ID string
+ Model string
+ Text string
+ Done bool
+ Usage *Usage
+ ToolCallDelta *ToolCallDelta
+}
+
+// ToolCall represents a function call from the model.
+type ToolCall struct {
+ ID string
+ Name string
+ Arguments string // JSON string
+}
+
+// ToolCallDelta represents a streaming chunk of a tool call.
+type ToolCallDelta struct {
+ Index int
+ ID string
+ Name string
+ Arguments string
+}
+
+// ============================================================
+// Models Endpoint Types
+// ============================================================
+
+// ModelInfo describes a single model available through the gateway.
+type ModelInfo struct {
+ ID string `json:"id"`
+ Provider string `json:"provider"`
+}
+
+// ModelsResponse is returned by GET /v1/models.
+type ModelsResponse struct {
+ Object string `json:"object"`
+ Data []ModelInfo `json:"data"`
+}
+
+// ============================================================
+// Validation
+// ============================================================
+
+// Validate performs basic structural validation.
+func (r *ResponseRequest) Validate() error {
+ if r == nil {
+ return errors.New("request is nil")
+ }
+ if r.Model == "" {
+ return errors.New("model is required")
+ }
+ if r.Input.String == nil && len(r.Input.Items) == 0 {
+ return errors.New("input is required")
+ }
+ return nil
+}
+
+// getStringField is a helper to safely extract string fields from a map
+func getStringField(m map[string]interface{}, key string) string {
+ if val, ok := m[key].(string); ok {
+ return val
+ }
+ return ""
+}
+
+
+
package auth
+
+import (
+ "context"
+ "crypto/rsa"
+ "encoding/base64"
+ "encoding/json"
+ "fmt"
+ "log/slog"
+ "math/big"
+ "net/http"
+ "strings"
+ "sync"
+ "time"
+
+ "github.com/golang-jwt/jwt/v5"
+)
+
+// Config holds OIDC authentication configuration.
+type Config struct {
+ Enabled bool `yaml:"enabled"`
+ Issuer string `yaml:"issuer"` // e.g., "https://accounts.google.com"
+ Audience string `yaml:"audience"` // e.g., your client ID
+}
+
+// Middleware provides JWT validation middleware.
+type Middleware struct {
+ cfg Config
+ keys map[string]*rsa.PublicKey
+ mu sync.RWMutex
+ client *http.Client
+ logger *slog.Logger
+}
+
+// New creates an authentication middleware.
+func New(cfg Config, logger *slog.Logger) (*Middleware, error) {
+ if !cfg.Enabled {
+ return &Middleware{cfg: cfg, logger: logger}, nil
+ }
+
+ if cfg.Issuer == "" {
+ return nil, fmt.Errorf("auth enabled but issuer not configured")
+ }
+
+ m := &Middleware{
+ cfg: cfg,
+ keys: make(map[string]*rsa.PublicKey),
+ client: &http.Client{Timeout: 10 * time.Second},
+ logger: logger,
+ }
+
+ // Fetch JWKS on startup
+ if err := m.refreshJWKS(); err != nil {
+ return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
+ }
+
+ // Refresh JWKS periodically
+ go m.periodicRefresh()
+
+ return m, nil
+}
+
+// Handler wraps an HTTP handler with authentication.
+func (m *Middleware) Handler(next http.Handler) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if !m.cfg.Enabled {
+ next.ServeHTTP(w, r)
+ return
+ }
+
+ // Extract token from Authorization header
+ authHeader := r.Header.Get("Authorization")
+ if authHeader == "" {
+ http.Error(w, "missing authorization header", http.StatusUnauthorized)
+ return
+ }
+
+ parts := strings.SplitN(authHeader, " ", 2)
+ if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
+ http.Error(w, "invalid authorization header format", http.StatusUnauthorized)
+ return
+ }
+
+ tokenString := parts[1]
+
+ // Validate token
+ claims, err := m.validateToken(tokenString)
+ if err != nil {
+ http.Error(w, fmt.Sprintf("invalid token: %v", err), http.StatusUnauthorized)
+ return
+ }
+
+ // Add claims to context
+ ctx := context.WithValue(r.Context(), claimsKey, claims)
+ next.ServeHTTP(w, r.WithContext(ctx))
+ })
+}
+
+type contextKey string
+
+const claimsKey contextKey = "jwt_claims"
+
+// GetClaims extracts JWT claims from request context.
+func GetClaims(ctx context.Context) (jwt.MapClaims, bool) {
+ claims, ok := ctx.Value(claimsKey).(jwt.MapClaims)
+ return claims, ok
+}
+
+func (m *Middleware) validateToken(tokenString string) (jwt.MapClaims, error) {
+ token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
+ // Verify signing method
+ if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
+ return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
+ }
+
+ // Get key ID from token header
+ kid, ok := token.Header["kid"].(string)
+ if !ok {
+ return nil, fmt.Errorf("missing kid in token header")
+ }
+
+ // Get public key
+ m.mu.RLock()
+ key, exists := m.keys[kid]
+ m.mu.RUnlock()
+
+ if !exists {
+ // Try refreshing JWKS
+ if err := m.refreshJWKS(); err != nil {
+ return nil, fmt.Errorf("failed to refresh JWKS: %w", err)
+ }
+
+ m.mu.RLock()
+ key, exists = m.keys[kid]
+ m.mu.RUnlock()
+
+ if !exists {
+ return nil, fmt.Errorf("unknown key ID: %s", kid)
+ }
+ }
+
+ return key, nil
+ })
+
+ if err != nil {
+ return nil, err
+ }
+
+ claims, ok := token.Claims.(jwt.MapClaims)
+ if !ok || !token.Valid {
+ return nil, fmt.Errorf("invalid token claims")
+ }
+
+ // Validate issuer
+ if iss, ok := claims["iss"].(string); !ok || iss != m.cfg.Issuer {
+ return nil, fmt.Errorf("invalid issuer: %s", iss)
+ }
+
+ // Validate audience if configured
+ if m.cfg.Audience != "" {
+ aud, ok := claims["aud"].(string)
+ if !ok {
+ // aud might be an array
+ audArray, ok := claims["aud"].([]interface{})
+ if !ok {
+ return nil, fmt.Errorf("invalid audience format")
+ }
+ found := false
+ for _, a := range audArray {
+ if audStr, ok := a.(string); ok && audStr == m.cfg.Audience {
+ found = true
+ break
+ }
+ }
+ if !found {
+ return nil, fmt.Errorf("audience not matched")
+ }
+ } else if aud != m.cfg.Audience {
+ return nil, fmt.Errorf("invalid audience: %s", aud)
+ }
+ }
+
+ return claims, nil
+}
+
+func (m *Middleware) refreshJWKS() error {
+ jwksURL := strings.TrimSuffix(m.cfg.Issuer, "/") + "/.well-known/openid-configuration"
+
+ // Fetch OIDC discovery document
+ resp, err := m.client.Get(jwksURL)
+ if err != nil {
+ return err
+ }
+ defer resp.Body.Close()
+
+ var oidcConfig struct {
+ JwksURI string `json:"jwks_uri"`
+ }
+ if err := json.NewDecoder(resp.Body).Decode(&oidcConfig); err != nil {
+ return err
+ }
+
+ // Fetch JWKS
+ resp, err = m.client.Get(oidcConfig.JwksURI)
+ if err != nil {
+ return err
+ }
+ defer resp.Body.Close()
+
+ var jwks struct {
+ Keys []struct {
+ Kid string `json:"kid"`
+ Kty string `json:"kty"`
+ Use string `json:"use"`
+ N string `json:"n"`
+ E string `json:"e"`
+ } `json:"keys"`
+ }
+
+ if err := json.NewDecoder(resp.Body).Decode(&jwks); err != nil {
+ return err
+ }
+
+ // Parse keys
+ newKeys := make(map[string]*rsa.PublicKey)
+ for _, key := range jwks.Keys {
+ if key.Kty != "RSA" || key.Use != "sig" {
+ continue
+ }
+
+ nBytes, err := base64.RawURLEncoding.DecodeString(key.N)
+ if err != nil {
+ continue
+ }
+
+ eBytes, err := base64.RawURLEncoding.DecodeString(key.E)
+ if err != nil {
+ continue
+ }
+
+ pubKey := &rsa.PublicKey{
+ N: new(big.Int).SetBytes(nBytes),
+ E: int(new(big.Int).SetBytes(eBytes).Int64()),
+ }
+
+ newKeys[key.Kid] = pubKey
+ }
+
+ m.mu.Lock()
+ m.keys = newKeys
+ m.mu.Unlock()
+
+ return nil
+}
+
+func (m *Middleware) periodicRefresh() {
+ ticker := time.NewTicker(1 * time.Hour)
+ defer ticker.Stop()
+
+ for range ticker.C {
+ if err := m.refreshJWKS(); err != nil {
+ m.logger.Error("failed to refresh JWKS",
+ slog.String("issuer", m.cfg.Issuer),
+ slog.String("error", err.Error()),
+ )
+ } else {
+ m.logger.Debug("successfully refreshed JWKS",
+ slog.String("issuer", m.cfg.Issuer),
+ )
+ }
+ }
+}
+
+
+
package config
+
+import (
+ "fmt"
+ "os"
+
+ "gopkg.in/yaml.v3"
+)
+
+// Config describes the full gateway configuration file.
+type Config struct {
+ Server ServerConfig `yaml:"server"`
+ Providers map[string]ProviderEntry `yaml:"providers"`
+ Models []ModelEntry `yaml:"models"`
+ Auth AuthConfig `yaml:"auth"`
+ Conversations ConversationConfig `yaml:"conversations"`
+ Logging LoggingConfig `yaml:"logging"`
+ RateLimit RateLimitConfig `yaml:"rate_limit"`
+ Observability ObservabilityConfig `yaml:"observability"`
+}
+
+// ConversationConfig controls conversation storage.
+type ConversationConfig struct {
+ // Store is the storage backend: "memory" (default), "sql", or "redis".
+ Store string `yaml:"store"`
+ // TTL is the conversation expiration duration (e.g. "1h", "30m"). Defaults to "1h".
+ TTL string `yaml:"ttl"`
+ // DSN is the database/Redis connection string, required when store is "sql" or "redis".
+ // Examples: "conversations.db" (SQLite), "postgres://user:pass@host/db", "redis://:password@localhost:6379/0".
+ DSN string `yaml:"dsn"`
+ // Driver is the SQL driver name, required when store is "sql".
+ // Examples: "sqlite3", "postgres", "mysql".
+ Driver string `yaml:"driver"`
+}
+
+// LoggingConfig controls logging format and level.
+type LoggingConfig struct {
+ // Format is the log output format: "json" (default) or "text".
+ Format string `yaml:"format"`
+ // Level is the minimum log level: "debug", "info" (default), "warn", or "error".
+ Level string `yaml:"level"`
+}
+
+// RateLimitConfig controls rate limiting behavior.
+type RateLimitConfig struct {
+ // Enabled controls whether rate limiting is active.
+ Enabled bool `yaml:"enabled"`
+ // RequestsPerSecond is the number of requests allowed per second per IP.
+ RequestsPerSecond float64 `yaml:"requests_per_second"`
+ // Burst is the maximum burst size allowed.
+ Burst int `yaml:"burst"`
+}
+
+// ObservabilityConfig controls observability features.
+type ObservabilityConfig struct {
+ Enabled bool `yaml:"enabled"`
+ Metrics MetricsConfig `yaml:"metrics"`
+ Tracing TracingConfig `yaml:"tracing"`
+}
+
+// MetricsConfig controls Prometheus metrics.
+type MetricsConfig struct {
+ Enabled bool `yaml:"enabled"`
+ Path string `yaml:"path"` // default: "/metrics"
+}
+
+// TracingConfig controls OpenTelemetry tracing.
+type TracingConfig struct {
+ Enabled bool `yaml:"enabled"`
+ ServiceName string `yaml:"service_name"` // default: "llm-gateway"
+ Sampler SamplerConfig `yaml:"sampler"`
+ Exporter ExporterConfig `yaml:"exporter"`
+}
+
+// SamplerConfig controls trace sampling.
+type SamplerConfig struct {
+ Type string `yaml:"type"` // "always", "never", "probability"
+ Rate float64 `yaml:"rate"` // 0.0 to 1.0
+}
+
+// ExporterConfig controls trace exporters.
+type ExporterConfig struct {
+ Type string `yaml:"type"` // "otlp", "stdout"
+ Endpoint string `yaml:"endpoint"`
+ Insecure bool `yaml:"insecure"`
+ Headers map[string]string `yaml:"headers"`
+}
+
+// AuthConfig holds OIDC authentication settings.
+type AuthConfig struct {
+ Enabled bool `yaml:"enabled"`
+ Issuer string `yaml:"issuer"`
+ Audience string `yaml:"audience"`
+}
+
+// ServerConfig controls HTTP server values.
+type ServerConfig struct {
+ Address string `yaml:"address"`
+ MaxRequestBodySize int64 `yaml:"max_request_body_size"` // Maximum request body size in bytes (default: 10MB)
+}
+
+// ProviderEntry defines a named provider instance in the config file.
+type ProviderEntry struct {
+ Type string `yaml:"type"`
+ APIKey string `yaml:"api_key"`
+ Endpoint string `yaml:"endpoint"`
+ APIVersion string `yaml:"api_version"`
+ Project string `yaml:"project"` // For Vertex AI
+ Location string `yaml:"location"` // For Vertex AI
+}
+
+// ModelEntry maps a model name to a provider entry.
+type ModelEntry struct {
+ Name string `yaml:"name"`
+ Provider string `yaml:"provider"`
+ ProviderModelID string `yaml:"provider_model_id"`
+}
+
+// ProviderConfig contains shared provider configuration fields used internally by providers.
+type ProviderConfig struct {
+ APIKey string `yaml:"api_key"`
+ Model string `yaml:"model"`
+ Endpoint string `yaml:"endpoint"`
+}
+
+// AzureOpenAIConfig contains Azure-specific settings used internally by the OpenAI provider.
+type AzureOpenAIConfig struct {
+ APIKey string `yaml:"api_key"`
+ Endpoint string `yaml:"endpoint"`
+ APIVersion string `yaml:"api_version"`
+}
+
+// AzureAnthropicConfig contains Azure-specific settings for Anthropic used internally.
+type AzureAnthropicConfig struct {
+ APIKey string `yaml:"api_key"`
+ Endpoint string `yaml:"endpoint"`
+ Model string `yaml:"model"`
+}
+
+// VertexAIConfig contains Vertex AI-specific settings used internally by the Google provider.
+type VertexAIConfig struct {
+ Project string `yaml:"project"`
+ Location string `yaml:"location"`
+}
+
+// Load reads and parses a YAML configuration file, expanding ${VAR} env references.
+func Load(path string) (*Config, error) {
+ data, err := os.ReadFile(path)
+ if err != nil {
+ return nil, fmt.Errorf("read config: %w", err)
+ }
+
+ expanded := os.Expand(string(data), os.Getenv)
+
+ var cfg Config
+ if err := yaml.Unmarshal([]byte(expanded), &cfg); err != nil {
+ return nil, fmt.Errorf("parse config: %w", err)
+ }
+
+ if err := cfg.validate(); err != nil {
+ return nil, err
+ }
+
+ return &cfg, nil
+}
+
+func (cfg *Config) validate() error {
+ for _, m := range cfg.Models {
+ if _, ok := cfg.Providers[m.Provider]; !ok {
+ return fmt.Errorf("model %q references unknown provider %q", m.Name, m.Provider)
+ }
+ }
+ return nil
+}
+
+
+
package conversation
+
+import (
+ "context"
+ "sync"
+ "time"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+)
+
+// Store defines the interface for conversation storage backends.
+type Store interface {
+ Get(ctx context.Context, id string) (*Conversation, error)
+ Create(ctx context.Context, id string, model string, messages []api.Message) (*Conversation, error)
+ Append(ctx context.Context, id string, messages ...api.Message) (*Conversation, error)
+ Delete(ctx context.Context, id string) error
+ Size() int
+ Close() error
+}
+
+// MemoryStore manages conversation history in-memory with automatic expiration.
+type MemoryStore struct {
+ conversations map[string]*Conversation
+ mu sync.RWMutex
+ ttl time.Duration
+ done chan struct{}
+}
+
+// Conversation holds the message history for a single conversation thread.
+type Conversation struct {
+ ID string
+ Messages []api.Message
+ Model string
+ CreatedAt time.Time
+ UpdatedAt time.Time
+}
+
+// NewMemoryStore creates an in-memory conversation store with the given TTL.
+func NewMemoryStore(ttl time.Duration) *MemoryStore {
+ s := &MemoryStore{
+ conversations: make(map[string]*Conversation),
+ ttl: ttl,
+ done: make(chan struct{}),
+ }
+
+ // Start cleanup goroutine if TTL is set
+ if ttl > 0 {
+ go s.cleanup()
+ }
+
+ return s
+}
+
+// Get retrieves a conversation by ID. Returns a deep copy to prevent data races.
+func (s *MemoryStore) Get(ctx context.Context, id string) (*Conversation, error) {
+ s.mu.RLock()
+ defer s.mu.RUnlock()
+
+ conv, ok := s.conversations[id]
+ if !ok {
+ return nil, nil
+ }
+
+ // Return a deep copy to prevent data races
+ msgsCopy := make([]api.Message, len(conv.Messages))
+ copy(msgsCopy, conv.Messages)
+
+ return &Conversation{
+ ID: conv.ID,
+ Messages: msgsCopy,
+ Model: conv.Model,
+ CreatedAt: conv.CreatedAt,
+ UpdatedAt: conv.UpdatedAt,
+ }, nil
+}
+
+// Create creates a new conversation with the given messages.
+func (s *MemoryStore) Create(ctx context.Context, id string, model string, messages []api.Message) (*Conversation, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+
+ now := time.Now()
+
+ // Store a copy to prevent external modifications
+ msgsCopy := make([]api.Message, len(messages))
+ copy(msgsCopy, messages)
+
+ conv := &Conversation{
+ ID: id,
+ Messages: msgsCopy,
+ Model: model,
+ CreatedAt: now,
+ UpdatedAt: now,
+ }
+
+ s.conversations[id] = conv
+
+ // Return a copy
+ return &Conversation{
+ ID: id,
+ Messages: messages,
+ Model: model,
+ CreatedAt: now,
+ UpdatedAt: now,
+ }, nil
+}
+
+// Append adds new messages to an existing conversation.
+func (s *MemoryStore) Append(ctx context.Context, id string, messages ...api.Message) (*Conversation, error) {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+
+ conv, ok := s.conversations[id]
+ if !ok {
+ return nil, nil
+ }
+
+ conv.Messages = append(conv.Messages, messages...)
+ conv.UpdatedAt = time.Now()
+
+ // Return a deep copy
+ msgsCopy := make([]api.Message, len(conv.Messages))
+ copy(msgsCopy, conv.Messages)
+
+ return &Conversation{
+ ID: conv.ID,
+ Messages: msgsCopy,
+ Model: conv.Model,
+ CreatedAt: conv.CreatedAt,
+ UpdatedAt: conv.UpdatedAt,
+ }, nil
+}
+
+// Delete removes a conversation from the store.
+func (s *MemoryStore) Delete(ctx context.Context, id string) error {
+ s.mu.Lock()
+ defer s.mu.Unlock()
+
+ delete(s.conversations, id)
+ return nil
+}
+
+// cleanup periodically removes expired conversations.
+func (s *MemoryStore) cleanup() {
+ ticker := time.NewTicker(1 * time.Minute)
+ defer ticker.Stop()
+
+ for {
+ select {
+ case <-ticker.C:
+ s.mu.Lock()
+ now := time.Now()
+ for id, conv := range s.conversations {
+ if now.Sub(conv.UpdatedAt) > s.ttl {
+ delete(s.conversations, id)
+ }
+ }
+ s.mu.Unlock()
+ case <-s.done:
+ return
+ }
+ }
+}
+
+// Size returns the number of active conversations.
+func (s *MemoryStore) Size() int {
+ s.mu.RLock()
+ defer s.mu.RUnlock()
+ return len(s.conversations)
+}
+
+// Close stops the cleanup goroutine and releases resources.
+func (s *MemoryStore) Close() error {
+ close(s.done)
+ return nil
+}
+
+
+
package conversation
+
+import (
+ "context"
+ "encoding/json"
+ "time"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/redis/go-redis/v9"
+)
+
+// RedisStore manages conversation history in Redis with automatic expiration.
+type RedisStore struct {
+ client *redis.Client
+ ttl time.Duration
+}
+
+// NewRedisStore creates a Redis-backed conversation store.
+func NewRedisStore(client *redis.Client, ttl time.Duration) *RedisStore {
+ return &RedisStore{
+ client: client,
+ ttl: ttl,
+ }
+}
+
+// key returns the Redis key for a conversation ID.
+func (s *RedisStore) key(id string) string {
+ return "conv:" + id
+}
+
+// Get retrieves a conversation by ID from Redis.
+func (s *RedisStore) Get(ctx context.Context, id string) (*Conversation, error) {
+ data, err := s.client.Get(ctx, s.key(id)).Bytes()
+ if err == redis.Nil {
+ return nil, nil
+ }
+ if err != nil {
+ return nil, err
+ }
+
+ var conv Conversation
+ if err := json.Unmarshal(data, &conv); err != nil {
+ return nil, err
+ }
+
+ return &conv, nil
+}
+
+// Create creates a new conversation with the given messages.
+func (s *RedisStore) Create(ctx context.Context, id string, model string, messages []api.Message) (*Conversation, error) {
+ now := time.Now()
+ conv := &Conversation{
+ ID: id,
+ Messages: messages,
+ Model: model,
+ CreatedAt: now,
+ UpdatedAt: now,
+ }
+
+ data, err := json.Marshal(conv)
+ if err != nil {
+ return nil, err
+ }
+
+ if err := s.client.Set(ctx, s.key(id), data, s.ttl).Err(); err != nil {
+ return nil, err
+ }
+
+ return conv, nil
+}
+
+// Append adds new messages to an existing conversation.
+func (s *RedisStore) Append(ctx context.Context, id string, messages ...api.Message) (*Conversation, error) {
+ conv, err := s.Get(ctx, id)
+ if err != nil {
+ return nil, err
+ }
+ if conv == nil {
+ return nil, nil
+ }
+
+ conv.Messages = append(conv.Messages, messages...)
+ conv.UpdatedAt = time.Now()
+
+ data, err := json.Marshal(conv)
+ if err != nil {
+ return nil, err
+ }
+
+ if err := s.client.Set(ctx, s.key(id), data, s.ttl).Err(); err != nil {
+ return nil, err
+ }
+
+ return conv, nil
+}
+
+// Delete removes a conversation from Redis.
+func (s *RedisStore) Delete(ctx context.Context, id string) error {
+ return s.client.Del(ctx, s.key(id)).Err()
+}
+
+// Size returns the number of active conversations in Redis.
+func (s *RedisStore) Size() int {
+ var count int
+ var cursor uint64
+ ctx := context.Background()
+
+ for {
+ keys, nextCursor, err := s.client.Scan(ctx, cursor, "conv:*", 100).Result()
+ if err != nil {
+ return 0
+ }
+
+ count += len(keys)
+ cursor = nextCursor
+
+ if cursor == 0 {
+ break
+ }
+ }
+
+ return count
+}
+
+// Close closes the Redis client connection.
+func (s *RedisStore) Close() error {
+ return s.client.Close()
+}
+
+
+
package conversation
+
+import (
+ "context"
+ "database/sql"
+ "encoding/json"
+ "time"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+)
+
+// sqlDialect holds driver-specific SQL statements.
+type sqlDialect struct {
+ getByID string
+ upsert string
+ update string
+ deleteByID string
+ cleanup string
+}
+
+func newDialect(driver string) sqlDialect {
+ if driver == "pgx" || driver == "postgres" {
+ return sqlDialect{
+ getByID: `SELECT id, model, messages, created_at, updated_at FROM conversations WHERE id = $1`,
+ upsert: `INSERT INTO conversations (id, model, messages, created_at, updated_at) VALUES ($1, $2, $3, $4, $5) ON CONFLICT (id) DO UPDATE SET model = EXCLUDED.model, messages = EXCLUDED.messages, updated_at = EXCLUDED.updated_at`,
+ update: `UPDATE conversations SET messages = $1, updated_at = $2 WHERE id = $3`,
+ deleteByID: `DELETE FROM conversations WHERE id = $1`,
+ cleanup: `DELETE FROM conversations WHERE updated_at < $1`,
+ }
+ }
+ return sqlDialect{
+ getByID: `SELECT id, model, messages, created_at, updated_at FROM conversations WHERE id = ?`,
+ upsert: `REPLACE INTO conversations (id, model, messages, created_at, updated_at) VALUES (?, ?, ?, ?, ?)`,
+ update: `UPDATE conversations SET messages = ?, updated_at = ? WHERE id = ?`,
+ deleteByID: `DELETE FROM conversations WHERE id = ?`,
+ cleanup: `DELETE FROM conversations WHERE updated_at < ?`,
+ }
+}
+
+// SQLStore manages conversation history in a SQL database with automatic expiration.
+type SQLStore struct {
+ db *sql.DB
+ ttl time.Duration
+ dialect sqlDialect
+ done chan struct{}
+}
+
+// NewSQLStore creates a SQL-backed conversation store. It creates the
+// conversations table if it does not already exist and starts a background
+// goroutine to remove expired rows.
+func NewSQLStore(db *sql.DB, driver string, ttl time.Duration) (*SQLStore, error) {
+ _, err := db.Exec(`CREATE TABLE IF NOT EXISTS conversations (
+ id TEXT PRIMARY KEY,
+ model TEXT NOT NULL,
+ messages TEXT NOT NULL,
+ created_at TIMESTAMP NOT NULL,
+ updated_at TIMESTAMP NOT NULL
+ )`)
+ if err != nil {
+ return nil, err
+ }
+
+ s := &SQLStore{
+ db: db,
+ ttl: ttl,
+ dialect: newDialect(driver),
+ done: make(chan struct{}),
+ }
+ if ttl > 0 {
+ go s.cleanup()
+ }
+ return s, nil
+}
+
+func (s *SQLStore) Get(ctx context.Context, id string) (*Conversation, error) {
+ row := s.db.QueryRowContext(ctx, s.dialect.getByID, id)
+
+ var conv Conversation
+ var msgJSON string
+ err := row.Scan(&conv.ID, &conv.Model, &msgJSON, &conv.CreatedAt, &conv.UpdatedAt)
+ if err == sql.ErrNoRows {
+ return nil, nil
+ }
+ if err != nil {
+ return nil, err
+ }
+
+ if err := json.Unmarshal([]byte(msgJSON), &conv.Messages); err != nil {
+ return nil, err
+ }
+
+ return &conv, nil
+}
+
+func (s *SQLStore) Create(ctx context.Context, id string, model string, messages []api.Message) (*Conversation, error) {
+ now := time.Now()
+ msgJSON, err := json.Marshal(messages)
+ if err != nil {
+ return nil, err
+ }
+
+ if _, err := s.db.ExecContext(ctx, s.dialect.upsert, id, model, string(msgJSON), now, now); err != nil {
+ return nil, err
+ }
+
+ return &Conversation{
+ ID: id,
+ Messages: messages,
+ Model: model,
+ CreatedAt: now,
+ UpdatedAt: now,
+ }, nil
+}
+
+func (s *SQLStore) Append(ctx context.Context, id string, messages ...api.Message) (*Conversation, error) {
+ conv, err := s.Get(ctx, id)
+ if err != nil {
+ return nil, err
+ }
+ if conv == nil {
+ return nil, nil
+ }
+
+ conv.Messages = append(conv.Messages, messages...)
+ conv.UpdatedAt = time.Now()
+
+ msgJSON, err := json.Marshal(conv.Messages)
+ if err != nil {
+ return nil, err
+ }
+
+ if _, err := s.db.ExecContext(ctx, s.dialect.update, string(msgJSON), conv.UpdatedAt, id); err != nil {
+ return nil, err
+ }
+
+ return conv, nil
+}
+
+func (s *SQLStore) Delete(ctx context.Context, id string) error {
+ _, err := s.db.ExecContext(ctx, s.dialect.deleteByID, id)
+ return err
+}
+
+func (s *SQLStore) Size() int {
+ var count int
+ _ = s.db.QueryRow(`SELECT COUNT(*) FROM conversations`).Scan(&count)
+ return count
+}
+
+func (s *SQLStore) cleanup() {
+ ticker := time.NewTicker(1 * time.Minute)
+ defer ticker.Stop()
+
+ for {
+ select {
+ case <-ticker.C:
+ cutoff := time.Now().Add(-s.ttl)
+ _, _ = s.db.Exec(s.dialect.cleanup, cutoff)
+ case <-s.done:
+ return
+ }
+ }
+}
+
+// Close stops the cleanup goroutine and closes the database connection.
+func (s *SQLStore) Close() error {
+ close(s.done)
+ return s.db.Close()
+}
+
+
+
package conversation
+
+import (
+ "context"
+ "database/sql"
+ "fmt"
+ "testing"
+ "time"
+
+ "github.com/alicebob/miniredis/v2"
+ _ "github.com/mattn/go-sqlite3"
+ "github.com/redis/go-redis/v9"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+)
+
+// SetupTestDB creates an in-memory SQLite database for testing
+func SetupTestDB(t *testing.T, driver string) *sql.DB {
+ t.Helper()
+
+ var dsn string
+ switch driver {
+ case "sqlite3":
+ // Use in-memory SQLite database
+ dsn = ":memory:"
+ case "postgres":
+ // For postgres tests, use a mock or skip
+ t.Skip("PostgreSQL tests require external database")
+ return nil
+ case "mysql":
+ // For mysql tests, use a mock or skip
+ t.Skip("MySQL tests require external database")
+ return nil
+ default:
+ t.Fatalf("unsupported driver: %s", driver)
+ return nil
+ }
+
+ db, err := sql.Open(driver, dsn)
+ if err != nil {
+ t.Fatalf("failed to open database: %v", err)
+ }
+
+ // Create the conversations table
+ schema := `
+ CREATE TABLE IF NOT EXISTS conversations (
+ conversation_id TEXT PRIMARY KEY,
+ messages TEXT NOT NULL,
+ updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+ )
+ `
+ if _, err := db.Exec(schema); err != nil {
+ db.Close()
+ t.Fatalf("failed to create schema: %v", err)
+ }
+
+ return db
+}
+
+// SetupTestRedis creates a miniredis instance for testing
+func SetupTestRedis(t *testing.T) (*redis.Client, *miniredis.Miniredis) {
+ t.Helper()
+
+ mr := miniredis.RunT(t)
+
+ client := redis.NewClient(&redis.Options{
+ Addr: mr.Addr(),
+ })
+
+ // Test connection
+ ctx := context.Background()
+ if err := client.Ping(ctx).Err(); err != nil {
+ t.Fatalf("failed to connect to miniredis: %v", err)
+ }
+
+ return client, mr
+}
+
+// CreateTestMessages generates test message fixtures
+func CreateTestMessages(count int) []api.Message {
+ messages := make([]api.Message, count)
+ for i := 0; i < count; i++ {
+ role := "user"
+ if i%2 == 1 {
+ role = "assistant"
+ }
+ messages[i] = api.Message{
+ Role: role,
+ Content: []api.ContentBlock{
+ {
+ Type: "text",
+ Text: fmt.Sprintf("Test message %d", i+1),
+ },
+ },
+ }
+ }
+ return messages
+}
+
+// CreateTestConversation creates a test conversation with the given ID and messages
+func CreateTestConversation(conversationID string, messageCount int) *Conversation {
+ return &Conversation{
+ ID: conversationID,
+ Messages: CreateTestMessages(messageCount),
+ Model: "test-model",
+ CreatedAt: time.Now(),
+ UpdatedAt: time.Now(),
+ }
+}
+
+// MockStore is a simple in-memory store for testing
+type MockStore struct {
+ conversations map[string]*Conversation
+ getCalled bool
+ createCalled bool
+ appendCalled bool
+ deleteCalled bool
+ sizeCalled bool
+}
+
+func NewMockStore() *MockStore {
+ return &MockStore{
+ conversations: make(map[string]*Conversation),
+ }
+}
+
+func (m *MockStore) Get(ctx context.Context, conversationID string) (*Conversation, error) {
+ m.getCalled = true
+ conv, ok := m.conversations[conversationID]
+ if !ok {
+ return nil, fmt.Errorf("conversation not found")
+ }
+ return conv, nil
+}
+
+func (m *MockStore) Create(ctx context.Context, conversationID string, model string, messages []api.Message) (*Conversation, error) {
+ m.createCalled = true
+ m.conversations[conversationID] = &Conversation{
+ ID: conversationID,
+ Model: model,
+ Messages: messages,
+ CreatedAt: time.Now(),
+ UpdatedAt: time.Now(),
+ }
+ return m.conversations[conversationID], nil
+}
+
+func (m *MockStore) Append(ctx context.Context, conversationID string, messages ...api.Message) (*Conversation, error) {
+ m.appendCalled = true
+ conv, ok := m.conversations[conversationID]
+ if !ok {
+ return nil, fmt.Errorf("conversation not found")
+ }
+ conv.Messages = append(conv.Messages, messages...)
+ conv.UpdatedAt = time.Now()
+ return conv, nil
+}
+
+func (m *MockStore) Delete(ctx context.Context, conversationID string) error {
+ m.deleteCalled = true
+ delete(m.conversations, conversationID)
+ return nil
+}
+
+func (m *MockStore) Size() int {
+ m.sizeCalled = true
+ return len(m.conversations)
+}
+
+func (m *MockStore) Close() error {
+ return nil
+}
+
+
+
package logger
+
+import (
+ "context"
+ "log/slog"
+ "os"
+
+ "go.opentelemetry.io/otel/trace"
+)
+
+type contextKey string
+
+const requestIDKey contextKey = "request_id"
+
+// New creates a logger with the specified format (json or text) and level.
+func New(format string, level string) *slog.Logger {
+ var handler slog.Handler
+
+ logLevel := parseLevel(level)
+ opts := &slog.HandlerOptions{
+ Level: logLevel,
+ AddSource: true, // Add file:line info for debugging
+ }
+
+ if format == "json" {
+ handler = slog.NewJSONHandler(os.Stdout, opts)
+ } else {
+ handler = slog.NewTextHandler(os.Stdout, opts)
+ }
+
+ return slog.New(handler)
+}
+
+// parseLevel converts a string level to slog.Level.
+func parseLevel(level string) slog.Level {
+ switch level {
+ case "debug":
+ return slog.LevelDebug
+ case "info":
+ return slog.LevelInfo
+ case "warn":
+ return slog.LevelWarn
+ case "error":
+ return slog.LevelError
+ default:
+ return slog.LevelInfo
+ }
+}
+
+// WithRequestID adds a request ID to the context for tracing.
+func WithRequestID(ctx context.Context, requestID string) context.Context {
+ return context.WithValue(ctx, requestIDKey, requestID)
+}
+
+// FromContext extracts the request ID from context, or returns empty string.
+func FromContext(ctx context.Context) string {
+ if id, ok := ctx.Value(requestIDKey).(string); ok {
+ return id
+ }
+ return ""
+}
+
+// LogAttrsWithTrace adds trace context to log attributes for correlation.
+func LogAttrsWithTrace(ctx context.Context, attrs ...any) []any {
+ spanCtx := trace.SpanFromContext(ctx).SpanContext()
+ if spanCtx.IsValid() {
+ attrs = append(attrs,
+ slog.String("trace_id", spanCtx.TraceID().String()),
+ slog.String("span_id", spanCtx.SpanID().String()),
+ )
+ }
+ return attrs
+}
+
+
+
package observability
+
+import (
+ "github.com/ajac-zero/latticelm/internal/conversation"
+ "github.com/ajac-zero/latticelm/internal/providers"
+ "github.com/prometheus/client_golang/prometheus"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+)
+
+// ProviderRegistry defines the interface for provider registries.
+// This matches the interface expected by the server.
+type ProviderRegistry interface {
+ Get(name string) (providers.Provider, bool)
+ Models() []struct{ Provider, Model string }
+ ResolveModelID(model string) string
+ Default(model string) (providers.Provider, error)
+}
+
+// WrapProviderRegistry wraps all providers in a registry with observability.
+func WrapProviderRegistry(registry ProviderRegistry, metricsRegistry *prometheus.Registry, tp *sdktrace.TracerProvider) ProviderRegistry {
+ if registry == nil {
+ return nil
+ }
+
+ // We can't directly modify the registry's internal map, so we'll need to
+ // wrap providers as they're retrieved. Instead, create a new instrumented registry.
+ return &InstrumentedRegistry{
+ base: registry,
+ metrics: metricsRegistry,
+ tracer: tp,
+ wrappedProviders: make(map[string]providers.Provider),
+ }
+}
+
+// InstrumentedRegistry wraps a provider registry to return instrumented providers.
+type InstrumentedRegistry struct {
+ base ProviderRegistry
+ metrics *prometheus.Registry
+ tracer *sdktrace.TracerProvider
+ wrappedProviders map[string]providers.Provider
+}
+
+// Get returns an instrumented provider by entry name.
+func (r *InstrumentedRegistry) Get(name string) (providers.Provider, bool) {
+ // Check if we've already wrapped this provider
+ if wrapped, ok := r.wrappedProviders[name]; ok {
+ return wrapped, true
+ }
+
+ // Get the base provider
+ p, ok := r.base.Get(name)
+ if !ok {
+ return nil, false
+ }
+
+ // Wrap it
+ wrapped := NewInstrumentedProvider(p, r.metrics, r.tracer)
+ r.wrappedProviders[name] = wrapped
+ return wrapped, true
+}
+
+// Default returns the instrumented provider for the given model name.
+func (r *InstrumentedRegistry) Default(model string) (providers.Provider, error) {
+ p, err := r.base.Default(model)
+ if err != nil {
+ return nil, err
+ }
+
+ // Check if we've already wrapped this provider
+ name := p.Name()
+ if wrapped, ok := r.wrappedProviders[name]; ok {
+ return wrapped, nil
+ }
+
+ // Wrap it
+ wrapped := NewInstrumentedProvider(p, r.metrics, r.tracer)
+ r.wrappedProviders[name] = wrapped
+ return wrapped, nil
+}
+
+// Models returns the list of configured models and their provider entry names.
+func (r *InstrumentedRegistry) Models() []struct{ Provider, Model string } {
+ return r.base.Models()
+}
+
+// ResolveModelID returns the provider_model_id for a model.
+func (r *InstrumentedRegistry) ResolveModelID(model string) string {
+ return r.base.ResolveModelID(model)
+}
+
+// WrapConversationStore wraps a conversation store with observability.
+func WrapConversationStore(store conversation.Store, backend string, metricsRegistry *prometheus.Registry, tp *sdktrace.TracerProvider) conversation.Store {
+ if store == nil {
+ return nil
+ }
+
+ return NewInstrumentedStore(store, backend, metricsRegistry, tp)
+}
+
+
+
package observability
+
+import (
+ "github.com/prometheus/client_golang/prometheus"
+)
+
+var (
+ // HTTP Metrics
+ httpRequestsTotal = prometheus.NewCounterVec(
+ prometheus.CounterOpts{
+ Name: "http_requests_total",
+ Help: "Total number of HTTP requests",
+ },
+ []string{"method", "path", "status"},
+ )
+
+ httpRequestDuration = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "http_request_duration_seconds",
+ Help: "HTTP request latency in seconds",
+ Buckets: []float64{0.01, 0.05, 0.1, 0.5, 1, 2.5, 5, 10, 30},
+ },
+ []string{"method", "path", "status"},
+ )
+
+ httpRequestSize = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "http_request_size_bytes",
+ Help: "HTTP request size in bytes",
+ Buckets: prometheus.ExponentialBuckets(100, 10, 7), // 100B to 100MB
+ },
+ []string{"method", "path"},
+ )
+
+ httpResponseSize = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "http_response_size_bytes",
+ Help: "HTTP response size in bytes",
+ Buckets: prometheus.ExponentialBuckets(100, 10, 7), // 100B to 100MB
+ },
+ []string{"method", "path"},
+ )
+
+ // Provider Metrics
+ providerRequestsTotal = prometheus.NewCounterVec(
+ prometheus.CounterOpts{
+ Name: "provider_requests_total",
+ Help: "Total number of provider requests",
+ },
+ []string{"provider", "model", "operation", "status"},
+ )
+
+ providerRequestDuration = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "provider_request_duration_seconds",
+ Help: "Provider request latency in seconds",
+ Buckets: []float64{0.1, 0.5, 1, 2, 5, 10, 20, 30, 60},
+ },
+ []string{"provider", "model", "operation"},
+ )
+
+ providerTokensTotal = prometheus.NewCounterVec(
+ prometheus.CounterOpts{
+ Name: "provider_tokens_total",
+ Help: "Total number of tokens processed",
+ },
+ []string{"provider", "model", "type"}, // type: input, output
+ )
+
+ providerStreamTTFB = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "provider_stream_ttfb_seconds",
+ Help: "Time to first byte for streaming requests in seconds",
+ Buckets: []float64{0.05, 0.1, 0.5, 1, 2, 5, 10},
+ },
+ []string{"provider", "model"},
+ )
+
+ providerStreamChunks = prometheus.NewCounterVec(
+ prometheus.CounterOpts{
+ Name: "provider_stream_chunks_total",
+ Help: "Total number of stream chunks received",
+ },
+ []string{"provider", "model"},
+ )
+
+ providerStreamDuration = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "provider_stream_duration_seconds",
+ Help: "Total duration of streaming requests in seconds",
+ Buckets: []float64{0.1, 0.5, 1, 2, 5, 10, 20, 30, 60},
+ },
+ []string{"provider", "model"},
+ )
+
+ // Conversation Store Metrics
+ conversationOperationsTotal = prometheus.NewCounterVec(
+ prometheus.CounterOpts{
+ Name: "conversation_operations_total",
+ Help: "Total number of conversation store operations",
+ },
+ []string{"operation", "backend", "status"},
+ )
+
+ conversationOperationDuration = prometheus.NewHistogramVec(
+ prometheus.HistogramOpts{
+ Name: "conversation_operation_duration_seconds",
+ Help: "Conversation store operation latency in seconds",
+ Buckets: []float64{0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1},
+ },
+ []string{"operation", "backend"},
+ )
+
+ conversationActiveCount = prometheus.NewGaugeVec(
+ prometheus.GaugeOpts{
+ Name: "conversation_active_count",
+ Help: "Number of active conversations",
+ },
+ []string{"backend"},
+ )
+
+ // Circuit Breaker Metrics
+ circuitBreakerState = prometheus.NewGaugeVec(
+ prometheus.GaugeOpts{
+ Name: "circuit_breaker_state",
+ Help: "Circuit breaker state (0=closed, 1=open, 2=half-open)",
+ },
+ []string{"provider"},
+ )
+
+ circuitBreakerStateTransitions = prometheus.NewCounterVec(
+ prometheus.CounterOpts{
+ Name: "circuit_breaker_state_transitions_total",
+ Help: "Total number of circuit breaker state transitions",
+ },
+ []string{"provider", "from", "to"},
+ )
+)
+
+// InitMetrics registers all metrics with a new Prometheus registry.
+func InitMetrics() *prometheus.Registry {
+ registry := prometheus.NewRegistry()
+
+ // Register HTTP metrics
+ registry.MustRegister(httpRequestsTotal)
+ registry.MustRegister(httpRequestDuration)
+ registry.MustRegister(httpRequestSize)
+ registry.MustRegister(httpResponseSize)
+
+ // Register provider metrics
+ registry.MustRegister(providerRequestsTotal)
+ registry.MustRegister(providerRequestDuration)
+ registry.MustRegister(providerTokensTotal)
+ registry.MustRegister(providerStreamTTFB)
+ registry.MustRegister(providerStreamChunks)
+ registry.MustRegister(providerStreamDuration)
+
+ // Register conversation store metrics
+ registry.MustRegister(conversationOperationsTotal)
+ registry.MustRegister(conversationOperationDuration)
+ registry.MustRegister(conversationActiveCount)
+
+ // Register circuit breaker metrics
+ registry.MustRegister(circuitBreakerState)
+ registry.MustRegister(circuitBreakerStateTransitions)
+
+ return registry
+}
+
+// RecordCircuitBreakerStateChange records a circuit breaker state transition.
+func RecordCircuitBreakerStateChange(provider, from, to string) {
+ // Record the transition
+ circuitBreakerStateTransitions.WithLabelValues(provider, from, to).Inc()
+
+ // Update the current state gauge
+ var stateValue float64
+ switch to {
+ case "closed":
+ stateValue = 0
+ case "open":
+ stateValue = 1
+ case "half-open":
+ stateValue = 2
+ }
+ circuitBreakerState.WithLabelValues(provider).Set(stateValue)
+}
+
+
+
package observability
+
+import (
+ "net/http"
+ "strconv"
+ "time"
+
+ "github.com/prometheus/client_golang/prometheus"
+)
+
+// MetricsMiddleware creates a middleware that records HTTP metrics.
+func MetricsMiddleware(next http.Handler, registry *prometheus.Registry, _ interface{}) http.Handler {
+ if registry == nil {
+ // If metrics are not enabled, pass through without modification
+ return next
+ }
+
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ start := time.Now()
+
+ // Record request size
+ if r.ContentLength > 0 {
+ httpRequestSize.WithLabelValues(r.Method, r.URL.Path).Observe(float64(r.ContentLength))
+ }
+
+ // Wrap response writer to capture status code and response size
+ wrapped := &metricsResponseWriter{
+ ResponseWriter: w,
+ statusCode: http.StatusOK,
+ bytesWritten: 0,
+ }
+
+ // Call the next handler
+ next.ServeHTTP(wrapped, r)
+
+ // Record metrics after request completes
+ duration := time.Since(start).Seconds()
+ status := strconv.Itoa(wrapped.statusCode)
+
+ httpRequestsTotal.WithLabelValues(r.Method, r.URL.Path, status).Inc()
+ httpRequestDuration.WithLabelValues(r.Method, r.URL.Path, status).Observe(duration)
+ httpResponseSize.WithLabelValues(r.Method, r.URL.Path).Observe(float64(wrapped.bytesWritten))
+ })
+}
+
+// metricsResponseWriter wraps http.ResponseWriter to capture status code and bytes written.
+type metricsResponseWriter struct {
+ http.ResponseWriter
+ statusCode int
+ bytesWritten int
+}
+
+func (w *metricsResponseWriter) WriteHeader(statusCode int) {
+ w.statusCode = statusCode
+ w.ResponseWriter.WriteHeader(statusCode)
+}
+
+func (w *metricsResponseWriter) Write(b []byte) (int, error) {
+ n, err := w.ResponseWriter.Write(b)
+ w.bytesWritten += n
+ return n, err
+}
+
+
+
package observability
+
+import (
+ "context"
+ "time"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/providers"
+ "github.com/prometheus/client_golang/prometheus"
+ "go.opentelemetry.io/otel/attribute"
+ "go.opentelemetry.io/otel/codes"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+ "go.opentelemetry.io/otel/trace"
+)
+
+// InstrumentedProvider wraps a provider with metrics and tracing.
+type InstrumentedProvider struct {
+ base providers.Provider
+ registry *prometheus.Registry
+ tracer trace.Tracer
+}
+
+// NewInstrumentedProvider wraps a provider with observability.
+func NewInstrumentedProvider(p providers.Provider, registry *prometheus.Registry, tp *sdktrace.TracerProvider) providers.Provider {
+ var tracer trace.Tracer
+ if tp != nil {
+ tracer = tp.Tracer("llm-gateway")
+ }
+
+ return &InstrumentedProvider{
+ base: p,
+ registry: registry,
+ tracer: tracer,
+ }
+}
+
+// Name returns the name of the underlying provider.
+func (p *InstrumentedProvider) Name() string {
+ return p.base.Name()
+}
+
+// Generate wraps the provider's Generate method with metrics and tracing.
+func (p *InstrumentedProvider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
+ // Start span if tracing is enabled
+ if p.tracer != nil {
+ var span trace.Span
+ ctx, span = p.tracer.Start(ctx, "provider.generate",
+ trace.WithSpanKind(trace.SpanKindClient),
+ trace.WithAttributes(
+ attribute.String("provider.name", p.base.Name()),
+ attribute.String("provider.model", req.Model),
+ ),
+ )
+ defer span.End()
+ }
+
+ // Record start time
+ start := time.Now()
+
+ // Call underlying provider
+ result, err := p.base.Generate(ctx, messages, req)
+
+ // Record metrics
+ duration := time.Since(start).Seconds()
+ status := "success"
+ if err != nil {
+ status = "error"
+ if p.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.RecordError(err)
+ span.SetStatus(codes.Error, err.Error())
+ }
+ } else if result != nil {
+ // Add token attributes to span
+ if p.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.SetAttributes(
+ attribute.Int64("provider.input_tokens", int64(result.Usage.InputTokens)),
+ attribute.Int64("provider.output_tokens", int64(result.Usage.OutputTokens)),
+ attribute.Int64("provider.total_tokens", int64(result.Usage.TotalTokens)),
+ )
+ span.SetStatus(codes.Ok, "")
+ }
+
+ // Record token metrics
+ if p.registry != nil {
+ providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "input").Add(float64(result.Usage.InputTokens))
+ providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "output").Add(float64(result.Usage.OutputTokens))
+ }
+ }
+
+ // Record request metrics
+ if p.registry != nil {
+ providerRequestsTotal.WithLabelValues(p.base.Name(), req.Model, "generate", status).Inc()
+ providerRequestDuration.WithLabelValues(p.base.Name(), req.Model, "generate").Observe(duration)
+ }
+
+ return result, err
+}
+
+// GenerateStream wraps the provider's GenerateStream method with metrics and tracing.
+func (p *InstrumentedProvider) GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error) {
+ // Start span if tracing is enabled
+ if p.tracer != nil {
+ var span trace.Span
+ ctx, span = p.tracer.Start(ctx, "provider.generate_stream",
+ trace.WithSpanKind(trace.SpanKindClient),
+ trace.WithAttributes(
+ attribute.String("provider.name", p.base.Name()),
+ attribute.String("provider.model", req.Model),
+ ),
+ )
+ defer span.End()
+ }
+
+ // Record start time
+ start := time.Now()
+ var ttfb time.Duration
+ firstChunk := true
+
+ // Create instrumented channels
+ baseChan, baseErrChan := p.base.GenerateStream(ctx, messages, req)
+ outChan := make(chan *api.ProviderStreamDelta)
+ outErrChan := make(chan error, 1)
+
+ // Metrics tracking
+ var chunkCount int64
+ var totalInputTokens, totalOutputTokens int64
+ var streamErr error
+
+ go func() {
+ defer close(outChan)
+ defer close(outErrChan)
+
+ for {
+ select {
+ case delta, ok := <-baseChan:
+ if !ok {
+ // Stream finished - record final metrics
+ duration := time.Since(start).Seconds()
+ status := "success"
+ if streamErr != nil {
+ status = "error"
+ if p.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.RecordError(streamErr)
+ span.SetStatus(codes.Error, streamErr.Error())
+ }
+ } else {
+ if p.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.SetAttributes(
+ attribute.Int64("provider.input_tokens", totalInputTokens),
+ attribute.Int64("provider.output_tokens", totalOutputTokens),
+ attribute.Int64("provider.chunk_count", chunkCount),
+ attribute.Float64("provider.ttfb_seconds", ttfb.Seconds()),
+ )
+ span.SetStatus(codes.Ok, "")
+ }
+
+ // Record token metrics
+ if p.registry != nil && (totalInputTokens > 0 || totalOutputTokens > 0) {
+ providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "input").Add(float64(totalInputTokens))
+ providerTokensTotal.WithLabelValues(p.base.Name(), req.Model, "output").Add(float64(totalOutputTokens))
+ }
+ }
+
+ // Record stream metrics
+ if p.registry != nil {
+ providerRequestsTotal.WithLabelValues(p.base.Name(), req.Model, "generate_stream", status).Inc()
+ providerStreamDuration.WithLabelValues(p.base.Name(), req.Model).Observe(duration)
+ providerStreamChunks.WithLabelValues(p.base.Name(), req.Model).Add(float64(chunkCount))
+ if ttfb > 0 {
+ providerStreamTTFB.WithLabelValues(p.base.Name(), req.Model).Observe(ttfb.Seconds())
+ }
+ }
+ return
+ }
+
+ // Record TTFB on first chunk
+ if firstChunk {
+ ttfb = time.Since(start)
+ firstChunk = false
+ }
+
+ chunkCount++
+
+ // Track token usage
+ if delta.Usage != nil {
+ totalInputTokens = int64(delta.Usage.InputTokens)
+ totalOutputTokens = int64(delta.Usage.OutputTokens)
+ }
+
+ // Forward the delta
+ outChan <- delta
+
+ case err, ok := <-baseErrChan:
+ if ok && err != nil {
+ streamErr = err
+ outErrChan <- err
+ }
+ return
+ }
+ }
+ }()
+
+ return outChan, outErrChan
+}
+
+
+
package observability
+
+import (
+ "context"
+ "time"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/conversation"
+ "github.com/prometheus/client_golang/prometheus"
+ "go.opentelemetry.io/otel/attribute"
+ "go.opentelemetry.io/otel/codes"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+ "go.opentelemetry.io/otel/trace"
+)
+
+// InstrumentedStore wraps a conversation store with metrics and tracing.
+type InstrumentedStore struct {
+ base conversation.Store
+ registry *prometheus.Registry
+ tracer trace.Tracer
+ backend string
+}
+
+// NewInstrumentedStore wraps a conversation store with observability.
+func NewInstrumentedStore(s conversation.Store, backend string, registry *prometheus.Registry, tp *sdktrace.TracerProvider) conversation.Store {
+ var tracer trace.Tracer
+ if tp != nil {
+ tracer = tp.Tracer("llm-gateway")
+ }
+
+ // Initialize gauge with current size
+ if registry != nil {
+ conversationActiveCount.WithLabelValues(backend).Set(float64(s.Size()))
+ }
+
+ return &InstrumentedStore{
+ base: s,
+ registry: registry,
+ tracer: tracer,
+ backend: backend,
+ }
+}
+
+// Get wraps the store's Get method with metrics and tracing.
+func (s *InstrumentedStore) Get(ctx context.Context, id string) (*conversation.Conversation, error) {
+ // Start span if tracing is enabled
+ if s.tracer != nil {
+ var span trace.Span
+ ctx, span = s.tracer.Start(ctx, "conversation.get",
+ trace.WithAttributes(
+ attribute.String("conversation.id", id),
+ attribute.String("conversation.backend", s.backend),
+ ),
+ )
+ defer span.End()
+ }
+
+ // Record start time
+ start := time.Now()
+
+ // Call underlying store
+ conv, err := s.base.Get(ctx, id)
+
+ // Record metrics
+ duration := time.Since(start).Seconds()
+ status := "success"
+ if err != nil {
+ status = "error"
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.RecordError(err)
+ span.SetStatus(codes.Error, err.Error())
+ }
+ } else {
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ if conv != nil {
+ span.SetAttributes(
+ attribute.Int("conversation.message_count", len(conv.Messages)),
+ attribute.String("conversation.model", conv.Model),
+ )
+ }
+ span.SetStatus(codes.Ok, "")
+ }
+ }
+
+ if s.registry != nil {
+ conversationOperationsTotal.WithLabelValues("get", s.backend, status).Inc()
+ conversationOperationDuration.WithLabelValues("get", s.backend).Observe(duration)
+ }
+
+ return conv, err
+}
+
+// Create wraps the store's Create method with metrics and tracing.
+func (s *InstrumentedStore) Create(ctx context.Context, id string, model string, messages []api.Message) (*conversation.Conversation, error) {
+ // Start span if tracing is enabled
+ if s.tracer != nil {
+ var span trace.Span
+ ctx, span = s.tracer.Start(ctx, "conversation.create",
+ trace.WithAttributes(
+ attribute.String("conversation.id", id),
+ attribute.String("conversation.backend", s.backend),
+ attribute.String("conversation.model", model),
+ attribute.Int("conversation.initial_messages", len(messages)),
+ ),
+ )
+ defer span.End()
+ }
+
+ // Record start time
+ start := time.Now()
+
+ // Call underlying store
+ conv, err := s.base.Create(ctx, id, model, messages)
+
+ // Record metrics
+ duration := time.Since(start).Seconds()
+ status := "success"
+ if err != nil {
+ status = "error"
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.RecordError(err)
+ span.SetStatus(codes.Error, err.Error())
+ }
+ } else {
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.SetStatus(codes.Ok, "")
+ }
+ }
+
+ if s.registry != nil {
+ conversationOperationsTotal.WithLabelValues("create", s.backend, status).Inc()
+ conversationOperationDuration.WithLabelValues("create", s.backend).Observe(duration)
+ // Update active count
+ conversationActiveCount.WithLabelValues(s.backend).Set(float64(s.base.Size()))
+ }
+
+ return conv, err
+}
+
+// Append wraps the store's Append method with metrics and tracing.
+func (s *InstrumentedStore) Append(ctx context.Context, id string, messages ...api.Message) (*conversation.Conversation, error) {
+ // Start span if tracing is enabled
+ if s.tracer != nil {
+ var span trace.Span
+ ctx, span = s.tracer.Start(ctx, "conversation.append",
+ trace.WithAttributes(
+ attribute.String("conversation.id", id),
+ attribute.String("conversation.backend", s.backend),
+ attribute.Int("conversation.appended_messages", len(messages)),
+ ),
+ )
+ defer span.End()
+ }
+
+ // Record start time
+ start := time.Now()
+
+ // Call underlying store
+ conv, err := s.base.Append(ctx, id, messages...)
+
+ // Record metrics
+ duration := time.Since(start).Seconds()
+ status := "success"
+ if err != nil {
+ status = "error"
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.RecordError(err)
+ span.SetStatus(codes.Error, err.Error())
+ }
+ } else {
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ if conv != nil {
+ span.SetAttributes(
+ attribute.Int("conversation.total_messages", len(conv.Messages)),
+ )
+ }
+ span.SetStatus(codes.Ok, "")
+ }
+ }
+
+ if s.registry != nil {
+ conversationOperationsTotal.WithLabelValues("append", s.backend, status).Inc()
+ conversationOperationDuration.WithLabelValues("append", s.backend).Observe(duration)
+ }
+
+ return conv, err
+}
+
+// Delete wraps the store's Delete method with metrics and tracing.
+func (s *InstrumentedStore) Delete(ctx context.Context, id string) error {
+ // Start span if tracing is enabled
+ if s.tracer != nil {
+ var span trace.Span
+ ctx, span = s.tracer.Start(ctx, "conversation.delete",
+ trace.WithAttributes(
+ attribute.String("conversation.id", id),
+ attribute.String("conversation.backend", s.backend),
+ ),
+ )
+ defer span.End()
+ }
+
+ // Record start time
+ start := time.Now()
+
+ // Call underlying store
+ err := s.base.Delete(ctx, id)
+
+ // Record metrics
+ duration := time.Since(start).Seconds()
+ status := "success"
+ if err != nil {
+ status = "error"
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.RecordError(err)
+ span.SetStatus(codes.Error, err.Error())
+ }
+ } else {
+ if s.tracer != nil {
+ span := trace.SpanFromContext(ctx)
+ span.SetStatus(codes.Ok, "")
+ }
+ }
+
+ if s.registry != nil {
+ conversationOperationsTotal.WithLabelValues("delete", s.backend, status).Inc()
+ conversationOperationDuration.WithLabelValues("delete", s.backend).Observe(duration)
+ // Update active count
+ conversationActiveCount.WithLabelValues(s.backend).Set(float64(s.base.Size()))
+ }
+
+ return err
+}
+
+// Size returns the size of the underlying store.
+func (s *InstrumentedStore) Size() int {
+ return s.base.Size()
+}
+
+// Close wraps the store's Close method.
+func (s *InstrumentedStore) Close() error {
+ return s.base.Close()
+}
+
+
+
package observability
+
+import (
+ "context"
+ "io"
+
+ "github.com/prometheus/client_golang/prometheus"
+ "github.com/prometheus/client_golang/prometheus/testutil"
+ "go.opentelemetry.io/otel"
+ "go.opentelemetry.io/otel/sdk/resource"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+ "go.opentelemetry.io/otel/sdk/trace/tracetest"
+ semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+)
+
+// NewTestRegistry creates a new isolated Prometheus registry for testing
+func NewTestRegistry() *prometheus.Registry {
+ return prometheus.NewRegistry()
+}
+
+// NewTestTracer creates a no-op tracer for testing
+func NewTestTracer() (*sdktrace.TracerProvider, *tracetest.InMemoryExporter) {
+ exporter := tracetest.NewInMemoryExporter()
+ res := resource.NewSchemaless(
+ semconv.ServiceNameKey.String("test-service"),
+ )
+ tp := sdktrace.NewTracerProvider(
+ sdktrace.WithSyncer(exporter),
+ sdktrace.WithResource(res),
+ )
+ otel.SetTracerProvider(tp)
+ return tp, exporter
+}
+
+// GetMetricValue extracts a metric value from a registry
+func GetMetricValue(registry *prometheus.Registry, metricName string) (float64, error) {
+ metrics, err := registry.Gather()
+ if err != nil {
+ return 0, err
+ }
+
+ for _, mf := range metrics {
+ if mf.GetName() == metricName {
+ if len(mf.GetMetric()) > 0 {
+ m := mf.GetMetric()[0]
+ if m.GetCounter() != nil {
+ return m.GetCounter().GetValue(), nil
+ }
+ if m.GetGauge() != nil {
+ return m.GetGauge().GetValue(), nil
+ }
+ if m.GetHistogram() != nil {
+ return float64(m.GetHistogram().GetSampleCount()), nil
+ }
+ }
+ }
+ }
+
+ return 0, nil
+}
+
+// CountMetricsWithName counts how many metrics match the given name
+func CountMetricsWithName(registry *prometheus.Registry, metricName string) (int, error) {
+ metrics, err := registry.Gather()
+ if err != nil {
+ return 0, err
+ }
+
+ for _, mf := range metrics {
+ if mf.GetName() == metricName {
+ return len(mf.GetMetric()), nil
+ }
+ }
+
+ return 0, nil
+}
+
+// GetCounterValue is a helper to get counter values using testutil
+func GetCounterValue(counter prometheus.Counter) float64 {
+ return testutil.ToFloat64(counter)
+}
+
+// NewNoOpTracerProvider creates a tracer provider that discards all spans
+func NewNoOpTracerProvider() *sdktrace.TracerProvider {
+ return sdktrace.NewTracerProvider(
+ sdktrace.WithSpanProcessor(sdktrace.NewSimpleSpanProcessor(&noOpExporter{})),
+ )
+}
+
+// noOpExporter is an exporter that discards all spans
+type noOpExporter struct{}
+
+func (e *noOpExporter) ExportSpans(context.Context, []sdktrace.ReadOnlySpan) error {
+ return nil
+}
+
+func (e *noOpExporter) Shutdown(context.Context) error {
+ return nil
+}
+
+// ShutdownTracer is a helper to safely shutdown a tracer provider
+func ShutdownTracer(tp *sdktrace.TracerProvider) error {
+ if tp != nil {
+ return tp.Shutdown(context.Background())
+ }
+ return nil
+}
+
+// NewTestExporter creates a test exporter that writes to the provided writer
+type TestExporter struct {
+ writer io.Writer
+}
+
+func (e *TestExporter) ExportSpans(ctx context.Context, spans []sdktrace.ReadOnlySpan) error {
+ return nil
+}
+
+func (e *TestExporter) Shutdown(ctx context.Context) error {
+ return nil
+}
+
+
+
package observability
+
+import (
+ "context"
+ "fmt"
+
+ "github.com/ajac-zero/latticelm/internal/config"
+ "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
+ "go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
+ "go.opentelemetry.io/otel/sdk/resource"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+ semconv "go.opentelemetry.io/otel/semconv/v1.24.0"
+ "google.golang.org/grpc"
+ "google.golang.org/grpc/credentials/insecure"
+)
+
+// InitTracer initializes the OpenTelemetry tracer provider.
+func InitTracer(cfg config.TracingConfig) (*sdktrace.TracerProvider, error) {
+ // Create resource with service information
+ res, err := resource.Merge(
+ resource.Default(),
+ resource.NewWithAttributes(
+ semconv.SchemaURL,
+ semconv.ServiceName(cfg.ServiceName),
+ ),
+ )
+ if err != nil {
+ return nil, fmt.Errorf("failed to create resource: %w", err)
+ }
+
+ // Create exporter
+ var exporter sdktrace.SpanExporter
+ switch cfg.Exporter.Type {
+ case "otlp":
+ exporter, err = createOTLPExporter(cfg.Exporter)
+ if err != nil {
+ return nil, fmt.Errorf("failed to create OTLP exporter: %w", err)
+ }
+ case "stdout":
+ exporter, err = stdouttrace.New(
+ stdouttrace.WithPrettyPrint(),
+ )
+ if err != nil {
+ return nil, fmt.Errorf("failed to create stdout exporter: %w", err)
+ }
+ default:
+ return nil, fmt.Errorf("unsupported exporter type: %s", cfg.Exporter.Type)
+ }
+
+ // Create sampler
+ sampler := createSampler(cfg.Sampler)
+
+ // Create tracer provider
+ tp := sdktrace.NewTracerProvider(
+ sdktrace.WithBatcher(exporter),
+ sdktrace.WithResource(res),
+ sdktrace.WithSampler(sampler),
+ )
+
+ return tp, nil
+}
+
+// createOTLPExporter creates an OTLP gRPC exporter.
+func createOTLPExporter(cfg config.ExporterConfig) (sdktrace.SpanExporter, error) {
+ opts := []otlptracegrpc.Option{
+ otlptracegrpc.WithEndpoint(cfg.Endpoint),
+ }
+
+ if cfg.Insecure {
+ opts = append(opts, otlptracegrpc.WithTLSCredentials(insecure.NewCredentials()))
+ }
+
+ if len(cfg.Headers) > 0 {
+ opts = append(opts, otlptracegrpc.WithHeaders(cfg.Headers))
+ }
+
+ // Add dial options to ensure connection
+ opts = append(opts, otlptracegrpc.WithDialOption(grpc.WithBlock()))
+
+ return otlptracegrpc.New(context.Background(), opts...)
+}
+
+// createSampler creates a sampler based on the configuration.
+func createSampler(cfg config.SamplerConfig) sdktrace.Sampler {
+ switch cfg.Type {
+ case "always":
+ return sdktrace.AlwaysSample()
+ case "never":
+ return sdktrace.NeverSample()
+ case "probability":
+ return sdktrace.TraceIDRatioBased(cfg.Rate)
+ default:
+ // Default to 10% sampling
+ return sdktrace.TraceIDRatioBased(0.1)
+ }
+}
+
+// Shutdown gracefully shuts down the tracer provider.
+func Shutdown(ctx context.Context, tp *sdktrace.TracerProvider) error {
+ if tp == nil {
+ return nil
+ }
+ return tp.Shutdown(ctx)
+}
+
+
+
package observability
+
+import (
+ "net/http"
+
+ "go.opentelemetry.io/otel"
+ "go.opentelemetry.io/otel/attribute"
+ "go.opentelemetry.io/otel/codes"
+ "go.opentelemetry.io/otel/propagation"
+ sdktrace "go.opentelemetry.io/otel/sdk/trace"
+ "go.opentelemetry.io/otel/trace"
+)
+
+// TracingMiddleware creates a middleware that adds OpenTelemetry tracing to HTTP requests.
+func TracingMiddleware(next http.Handler, tp *sdktrace.TracerProvider) http.Handler {
+ if tp == nil {
+ // If tracing is not enabled, pass through without modification
+ return next
+ }
+
+ // Set up W3C Trace Context propagation
+ otel.SetTextMapPropagator(propagation.TraceContext{})
+
+ tracer := tp.Tracer("llm-gateway")
+
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ // Extract trace context from incoming request headers
+ ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+
+ // Start a new span
+ ctx, span := tracer.Start(ctx, "HTTP "+r.Method+" "+r.URL.Path,
+ trace.WithSpanKind(trace.SpanKindServer),
+ trace.WithAttributes(
+ attribute.String("http.method", r.Method),
+ attribute.String("http.route", r.URL.Path),
+ attribute.String("http.scheme", r.URL.Scheme),
+ attribute.String("http.host", r.Host),
+ attribute.String("http.user_agent", r.Header.Get("User-Agent")),
+ ),
+ )
+ defer span.End()
+
+ // Add request ID to span if present
+ if requestID := r.Header.Get("X-Request-ID"); requestID != "" {
+ span.SetAttributes(attribute.String("http.request_id", requestID))
+ }
+
+ // Create a response writer wrapper to capture status code
+ wrapped := &statusResponseWriter{
+ ResponseWriter: w,
+ statusCode: http.StatusOK,
+ }
+
+ // Inject trace context into request for downstream services
+ r = r.WithContext(ctx)
+
+ // Call the next handler
+ next.ServeHTTP(wrapped, r)
+
+ // Record the status code in the span
+ span.SetAttributes(attribute.Int("http.status_code", wrapped.statusCode))
+
+ // Set span status based on HTTP status code
+ if wrapped.statusCode >= 400 {
+ span.SetStatus(codes.Error, http.StatusText(wrapped.statusCode))
+ } else {
+ span.SetStatus(codes.Ok, "")
+ }
+ })
+}
+
+// statusResponseWriter wraps http.ResponseWriter to capture the status code.
+type statusResponseWriter struct {
+ http.ResponseWriter
+ statusCode int
+}
+
+func (w *statusResponseWriter) WriteHeader(statusCode int) {
+ w.statusCode = statusCode
+ w.ResponseWriter.WriteHeader(statusCode)
+}
+
+func (w *statusResponseWriter) Write(b []byte) (int, error) {
+ return w.ResponseWriter.Write(b)
+}
+
+
+
package anthropic
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+
+ "github.com/anthropics/anthropic-sdk-go"
+ "github.com/anthropics/anthropic-sdk-go/option"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/config"
+)
+
+const Name = "anthropic"
+
+// Provider implements the Anthropic SDK integration.
+// It supports both direct Anthropic API and Azure-hosted (Microsoft Foundry) endpoints.
+type Provider struct {
+ cfg config.ProviderConfig
+ client *anthropic.Client
+ azure bool
+}
+
+// New constructs a Provider for the direct Anthropic API.
+func New(cfg config.ProviderConfig) *Provider {
+ var client *anthropic.Client
+ if cfg.APIKey != "" {
+ c := anthropic.NewClient(option.WithAPIKey(cfg.APIKey))
+ client = &c
+ }
+ return &Provider{
+ cfg: cfg,
+ client: client,
+ }
+}
+
+// NewAzure constructs a Provider targeting Azure-hosted Anthropic (Microsoft Foundry).
+// The Azure endpoint uses api-key header auth and a base URL like
+// https://<resource>.services.ai.azure.com/anthropic.
+func NewAzure(azureCfg config.AzureAnthropicConfig) *Provider {
+ var client *anthropic.Client
+ if azureCfg.APIKey != "" && azureCfg.Endpoint != "" {
+ c := anthropic.NewClient(
+ option.WithBaseURL(azureCfg.Endpoint),
+ option.WithAPIKey("unused"),
+ option.WithAuthToken(azureCfg.APIKey),
+ )
+ client = &c
+ }
+ return &Provider{
+ cfg: config.ProviderConfig{
+ APIKey: azureCfg.APIKey,
+ Model: azureCfg.Model,
+ },
+ client: client,
+ azure: true,
+ }
+}
+
+func (p *Provider) Name() string { return Name }
+
+// Generate routes the request to Anthropic's API.
+func (p *Provider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
+ if p.cfg.APIKey == "" {
+ return nil, fmt.Errorf("anthropic api key missing")
+ }
+ if p.client == nil {
+ return nil, fmt.Errorf("anthropic client not initialized")
+ }
+
+ // Convert messages to Anthropic format
+ anthropicMsgs := make([]anthropic.MessageParam, 0, len(messages))
+ var system string
+
+ for _, msg := range messages {
+ var content string
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ content += block.Text
+ }
+ }
+
+ switch msg.Role {
+ case "user":
+ anthropicMsgs = append(anthropicMsgs, anthropic.NewUserMessage(anthropic.NewTextBlock(content)))
+ case "assistant":
+ // Build content blocks including text and tool calls
+ var contentBlocks []anthropic.ContentBlockParamUnion
+ if content != "" {
+ contentBlocks = append(contentBlocks, anthropic.NewTextBlock(content))
+ }
+ // Add tool use blocks
+ for _, tc := range msg.ToolCalls {
+ var input map[string]interface{}
+ if err := json.Unmarshal([]byte(tc.Arguments), &input); err != nil {
+ // If unmarshal fails, skip this tool call
+ continue
+ }
+ contentBlocks = append(contentBlocks, anthropic.NewToolUseBlock(tc.ID, input, tc.Name))
+ }
+ if len(contentBlocks) > 0 {
+ anthropicMsgs = append(anthropicMsgs, anthropic.NewAssistantMessage(contentBlocks...))
+ }
+ case "tool":
+ // Tool results must be in user message with tool_result blocks
+ anthropicMsgs = append(anthropicMsgs, anthropic.NewUserMessage(
+ anthropic.NewToolResultBlock(msg.CallID, content, false),
+ ))
+ case "system", "developer":
+ system = content
+ }
+ }
+
+ // Build request params
+ maxTokens := int64(4096)
+ if req.MaxOutputTokens != nil {
+ maxTokens = int64(*req.MaxOutputTokens)
+ }
+
+ params := anthropic.MessageNewParams{
+ Model: anthropic.Model(req.Model),
+ Messages: anthropicMsgs,
+ MaxTokens: maxTokens,
+ }
+
+ if system != "" {
+ systemBlocks := []anthropic.TextBlockParam{
+ {Text: system, Type: "text"},
+ }
+ params.System = systemBlocks
+ }
+
+ if req.Temperature != nil {
+ params.Temperature = anthropic.Float(*req.Temperature)
+ }
+ if req.TopP != nil {
+ params.TopP = anthropic.Float(*req.TopP)
+ }
+
+ // Add tools if present
+ if req.Tools != nil && len(req.Tools) > 0 {
+ tools, err := parseTools(req)
+ if err != nil {
+ return nil, fmt.Errorf("parse tools: %w", err)
+ }
+ params.Tools = tools
+ }
+
+ // Add tool_choice if present
+ if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
+ toolChoice, err := parseToolChoice(req)
+ if err != nil {
+ return nil, fmt.Errorf("parse tool_choice: %w", err)
+ }
+ params.ToolChoice = toolChoice
+ }
+
+ // Call Anthropic API
+ resp, err := p.client.Messages.New(ctx, params)
+ if err != nil {
+ return nil, fmt.Errorf("anthropic api error: %w", err)
+ }
+
+ // Extract text and tool calls from response
+ var text string
+ var toolCalls []api.ToolCall
+
+ for _, block := range resp.Content {
+ switch block.Type {
+ case "text":
+ text += block.AsText().Text
+ case "tool_use":
+ // Extract tool calls
+ toolUse := block.AsToolUse()
+ argsJSON, _ := json.Marshal(toolUse.Input)
+ toolCalls = append(toolCalls, api.ToolCall{
+ ID: toolUse.ID,
+ Name: toolUse.Name,
+ Arguments: string(argsJSON),
+ })
+ }
+ }
+
+ return &api.ProviderResult{
+ ID: resp.ID,
+ Model: string(resp.Model),
+ Text: text,
+ ToolCalls: toolCalls,
+ Usage: api.Usage{
+ InputTokens: int(resp.Usage.InputTokens),
+ OutputTokens: int(resp.Usage.OutputTokens),
+ TotalTokens: int(resp.Usage.InputTokens + resp.Usage.OutputTokens),
+ },
+ }, nil
+}
+
+// GenerateStream handles streaming requests to Anthropic.
+func (p *Provider) GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error) {
+ deltaChan := make(chan *api.ProviderStreamDelta)
+ errChan := make(chan error, 1)
+
+ go func() {
+ defer close(deltaChan)
+ defer close(errChan)
+
+ if p.cfg.APIKey == "" {
+ errChan <- fmt.Errorf("anthropic api key missing")
+ return
+ }
+ if p.client == nil {
+ errChan <- fmt.Errorf("anthropic client not initialized")
+ return
+ }
+
+ // Convert messages to Anthropic format
+ anthropicMsgs := make([]anthropic.MessageParam, 0, len(messages))
+ var system string
+
+ for _, msg := range messages {
+ var content string
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ content += block.Text
+ }
+ }
+
+ switch msg.Role {
+ case "user":
+ anthropicMsgs = append(anthropicMsgs, anthropic.NewUserMessage(anthropic.NewTextBlock(content)))
+ case "assistant":
+ // Build content blocks including text and tool calls
+ var contentBlocks []anthropic.ContentBlockParamUnion
+ if content != "" {
+ contentBlocks = append(contentBlocks, anthropic.NewTextBlock(content))
+ }
+ // Add tool use blocks
+ for _, tc := range msg.ToolCalls {
+ var input map[string]interface{}
+ if err := json.Unmarshal([]byte(tc.Arguments), &input); err != nil {
+ // If unmarshal fails, skip this tool call
+ continue
+ }
+ contentBlocks = append(contentBlocks, anthropic.NewToolUseBlock(tc.ID, input, tc.Name))
+ }
+ if len(contentBlocks) > 0 {
+ anthropicMsgs = append(anthropicMsgs, anthropic.NewAssistantMessage(contentBlocks...))
+ }
+ case "tool":
+ // Tool results must be in user message with tool_result blocks
+ anthropicMsgs = append(anthropicMsgs, anthropic.NewUserMessage(
+ anthropic.NewToolResultBlock(msg.CallID, content, false),
+ ))
+ case "system", "developer":
+ system = content
+ }
+ }
+
+ // Build params
+ maxTokens := int64(4096)
+ if req.MaxOutputTokens != nil {
+ maxTokens = int64(*req.MaxOutputTokens)
+ }
+
+ params := anthropic.MessageNewParams{
+ Model: anthropic.Model(req.Model),
+ Messages: anthropicMsgs,
+ MaxTokens: maxTokens,
+ }
+
+ if system != "" {
+ systemBlocks := []anthropic.TextBlockParam{
+ {Text: system, Type: "text"},
+ }
+ params.System = systemBlocks
+ }
+
+ if req.Temperature != nil {
+ params.Temperature = anthropic.Float(*req.Temperature)
+ }
+ if req.TopP != nil {
+ params.TopP = anthropic.Float(*req.TopP)
+ }
+
+ // Add tools if present
+ if req.Tools != nil && len(req.Tools) > 0 {
+ tools, err := parseTools(req)
+ if err != nil {
+ errChan <- fmt.Errorf("parse tools: %w", err)
+ return
+ }
+ params.Tools = tools
+ }
+
+ // Add tool_choice if present
+ if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
+ toolChoice, err := parseToolChoice(req)
+ if err != nil {
+ errChan <- fmt.Errorf("parse tool_choice: %w", err)
+ return
+ }
+ params.ToolChoice = toolChoice
+ }
+
+ // Create stream
+ stream := p.client.Messages.NewStreaming(ctx, params)
+
+ // Track content block index and tool call state
+ var contentBlockIndex int
+
+ // Process stream
+ for stream.Next() {
+ event := stream.Current()
+
+ switch event.Type {
+ case "content_block_start":
+ // New content block (text or tool_use)
+ contentBlockIndex = int(event.Index)
+ if event.ContentBlock.Type == "tool_use" {
+ // Send tool call delta with ID and name
+ toolUse := event.ContentBlock.AsToolUse()
+ delta := &api.ToolCallDelta{
+ Index: contentBlockIndex,
+ ID: toolUse.ID,
+ Name: toolUse.Name,
+ }
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{ToolCallDelta: delta}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ }
+
+ case "content_block_delta":
+ if event.Delta.Type == "text_delta" {
+ // Text streaming
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{Text: event.Delta.Text}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ } else if event.Delta.Type == "input_json_delta" {
+ // Tool arguments streaming
+ delta := &api.ToolCallDelta{
+ Index: int(event.Index),
+ Arguments: event.Delta.PartialJSON,
+ }
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{ToolCallDelta: delta}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ }
+ }
+ }
+
+ if err := stream.Err(); err != nil {
+ errChan <- fmt.Errorf("anthropic stream error: %w", err)
+ return
+ }
+
+ // Send final delta
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{Done: true}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ }
+ }()
+
+ return deltaChan, errChan
+}
+
+func chooseModel(requested, defaultModel string) string {
+ if requested != "" {
+ return requested
+ }
+ if defaultModel != "" {
+ return defaultModel
+ }
+ return "claude-3-5-sonnet"
+}
+
+
+
package anthropic
+
+import (
+ "encoding/json"
+ "fmt"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/anthropics/anthropic-sdk-go"
+)
+
+// parseTools converts Open Responses tools to Anthropic format
+func parseTools(req *api.ResponseRequest) ([]anthropic.ToolUnionParam, error) {
+ if req.Tools == nil || len(req.Tools) == 0 {
+ return nil, nil
+ }
+
+ var toolDefs []map[string]interface{}
+ if err := json.Unmarshal(req.Tools, &toolDefs); err != nil {
+ return nil, fmt.Errorf("unmarshal tools: %w", err)
+ }
+
+ var tools []anthropic.ToolUnionParam
+ for _, td := range toolDefs {
+ // Extract: name, description, parameters
+ // Note: Anthropic uses "input_schema" instead of "parameters"
+ name, _ := td["name"].(string)
+ desc, _ := td["description"].(string)
+ params, _ := td["parameters"].(map[string]interface{})
+
+ inputSchema := anthropic.ToolInputSchemaParam{
+ Type: "object",
+ Properties: params["properties"],
+ }
+
+ // Add required fields if present
+ if required, ok := params["required"].([]interface{}); ok {
+ requiredStrs := make([]string, 0, len(required))
+ for _, r := range required {
+ if str, ok := r.(string); ok {
+ requiredStrs = append(requiredStrs, str)
+ }
+ }
+ inputSchema.Required = requiredStrs
+ }
+
+ // Create the tool using ToolUnionParamOfTool
+ tool := anthropic.ToolUnionParamOfTool(inputSchema, name)
+
+ if desc != "" {
+ tool.OfTool.Description = anthropic.String(desc)
+ }
+
+ tools = append(tools, tool)
+ }
+
+ return tools, nil
+}
+
+// parseToolChoice converts Open Responses tool_choice to Anthropic format
+func parseToolChoice(req *api.ResponseRequest) (anthropic.ToolChoiceUnionParam, error) {
+ var result anthropic.ToolChoiceUnionParam
+
+ if req.ToolChoice == nil || len(req.ToolChoice) == 0 {
+ return result, nil
+ }
+
+ var choice interface{}
+ if err := json.Unmarshal(req.ToolChoice, &choice); err != nil {
+ return result, fmt.Errorf("unmarshal tool_choice: %w", err)
+ }
+
+ // Handle string values: "auto", "any", "required"
+ if str, ok := choice.(string); ok {
+ switch str {
+ case "auto":
+ result.OfAuto = &anthropic.ToolChoiceAutoParam{
+ Type: "auto",
+ }
+ case "any", "required":
+ result.OfAny = &anthropic.ToolChoiceAnyParam{
+ Type: "any",
+ }
+ case "none":
+ result.OfNone = &anthropic.ToolChoiceNoneParam{
+ Type: "none",
+ }
+ default:
+ return result, fmt.Errorf("unknown tool_choice string: %s", str)
+ }
+ return result, nil
+ }
+
+ // Handle specific tool selection: {"type": "tool", "function": {"name": "..."}}
+ if obj, ok := choice.(map[string]interface{}); ok {
+ // Check for OpenAI format: {"type": "function", "function": {"name": "..."}}
+ if funcObj, ok := obj["function"].(map[string]interface{}); ok {
+ if name, ok := funcObj["name"].(string); ok {
+ result.OfTool = &anthropic.ToolChoiceToolParam{
+ Type: "tool",
+ Name: name,
+ }
+ return result, nil
+ }
+ }
+
+ // Check for direct name field
+ if name, ok := obj["name"].(string); ok {
+ result.OfTool = &anthropic.ToolChoiceToolParam{
+ Type: "tool",
+ Name: name,
+ }
+ return result, nil
+ }
+ }
+
+ return result, fmt.Errorf("invalid tool_choice format")
+}
+
+// extractToolCalls converts Anthropic content blocks to api.ToolCall
+func extractToolCalls(content []anthropic.ContentBlockUnion) []api.ToolCall {
+ var toolCalls []api.ToolCall
+
+ for _, block := range content {
+ // Check if this is a tool_use block
+ if block.Type == "tool_use" {
+ // Cast to ToolUseBlock to access the fields
+ toolUse := block.AsToolUse()
+
+ // Marshal the input to JSON string for Arguments
+ argsJSON, _ := json.Marshal(toolUse.Input)
+
+ toolCalls = append(toolCalls, api.ToolCall{
+ ID: toolUse.ID,
+ Name: toolUse.Name,
+ Arguments: string(argsJSON),
+ })
+ }
+ }
+
+ return toolCalls
+}
+
+// extractToolCallDelta extracts tool call delta from streaming content block delta
+func extractToolCallDelta(delta anthropic.RawContentBlockDeltaUnion, index int) *api.ToolCallDelta {
+ // Check if this is an input_json_delta (streaming tool arguments)
+ if delta.Type == "input_json_delta" {
+ return &api.ToolCallDelta{
+ Index: index,
+ Arguments: delta.PartialJSON,
+ }
+ }
+
+ return nil
+}
+
+
+
package providers
+
+import (
+ "context"
+ "fmt"
+ "time"
+
+ "github.com/sony/gobreaker"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+)
+
+// CircuitBreakerProvider wraps a Provider with circuit breaker functionality.
+type CircuitBreakerProvider struct {
+ provider Provider
+ cb *gobreaker.CircuitBreaker
+}
+
+// CircuitBreakerConfig holds configuration for the circuit breaker.
+type CircuitBreakerConfig struct {
+ // MaxRequests is the maximum number of requests allowed to pass through
+ // when the circuit breaker is half-open. Default: 3
+ MaxRequests uint32
+
+ // Interval is the cyclic period of the closed state for the circuit breaker
+ // to clear the internal Counts. Default: 30s
+ Interval time.Duration
+
+ // Timeout is the period of the open state, after which the state becomes half-open.
+ // Default: 60s
+ Timeout time.Duration
+
+ // MinRequests is the minimum number of requests needed before evaluating failure ratio.
+ // Default: 5
+ MinRequests uint32
+
+ // FailureRatio is the ratio of failures that will trip the circuit breaker.
+ // Default: 0.5 (50%)
+ FailureRatio float64
+
+ // OnStateChange is an optional callback invoked when circuit breaker state changes.
+ // Parameters: provider name, from state, to state
+ OnStateChange func(provider, from, to string)
+}
+
+// DefaultCircuitBreakerConfig returns a sensible default configuration.
+func DefaultCircuitBreakerConfig() CircuitBreakerConfig {
+ return CircuitBreakerConfig{
+ MaxRequests: 3,
+ Interval: 30 * time.Second,
+ Timeout: 60 * time.Second,
+ MinRequests: 5,
+ FailureRatio: 0.5,
+ }
+}
+
+// NewCircuitBreakerProvider wraps a provider with circuit breaker functionality.
+func NewCircuitBreakerProvider(provider Provider, cfg CircuitBreakerConfig) *CircuitBreakerProvider {
+ providerName := provider.Name()
+
+ settings := gobreaker.Settings{
+ Name: fmt.Sprintf("%s-circuit-breaker", providerName),
+ MaxRequests: cfg.MaxRequests,
+ Interval: cfg.Interval,
+ Timeout: cfg.Timeout,
+ ReadyToTrip: func(counts gobreaker.Counts) bool {
+ // Only trip if we have enough requests to be statistically meaningful
+ if counts.Requests < cfg.MinRequests {
+ return false
+ }
+ failureRatio := float64(counts.TotalFailures) / float64(counts.Requests)
+ return failureRatio >= cfg.FailureRatio
+ },
+ OnStateChange: func(name string, from gobreaker.State, to gobreaker.State) {
+ // Call the callback if provided
+ if cfg.OnStateChange != nil {
+ cfg.OnStateChange(providerName, from.String(), to.String())
+ }
+ },
+ }
+
+ return &CircuitBreakerProvider{
+ provider: provider,
+ cb: gobreaker.NewCircuitBreaker(settings),
+ }
+}
+
+// Name returns the underlying provider name.
+func (p *CircuitBreakerProvider) Name() string {
+ return p.provider.Name()
+}
+
+// Generate wraps the provider's Generate method with circuit breaker protection.
+func (p *CircuitBreakerProvider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
+ result, err := p.cb.Execute(func() (interface{}, error) {
+ return p.provider.Generate(ctx, messages, req)
+ })
+
+ if err != nil {
+ return nil, err
+ }
+
+ return result.(*api.ProviderResult), nil
+}
+
+// GenerateStream wraps the provider's GenerateStream method with circuit breaker protection.
+func (p *CircuitBreakerProvider) GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error) {
+ // For streaming, we check the circuit breaker state before initiating the stream
+ // If the circuit is open, we return an error immediately
+ state := p.cb.State()
+ if state == gobreaker.StateOpen {
+ errChan := make(chan error, 1)
+ deltaChan := make(chan *api.ProviderStreamDelta)
+ errChan <- gobreaker.ErrOpenState
+ close(deltaChan)
+ close(errChan)
+ return deltaChan, errChan
+ }
+
+ // If circuit is closed or half-open, attempt the stream
+ deltaChan, errChan := p.provider.GenerateStream(ctx, messages, req)
+
+ // Wrap the error channel to report successes/failures to circuit breaker
+ wrappedErrChan := make(chan error, 1)
+
+ go func() {
+ defer close(wrappedErrChan)
+
+ // Wait for the error channel to signal completion
+ if err := <-errChan; err != nil {
+ // Record failure in circuit breaker
+ p.cb.Execute(func() (interface{}, error) {
+ return nil, err
+ })
+ wrappedErrChan <- err
+ } else {
+ // Record success in circuit breaker
+ p.cb.Execute(func() (interface{}, error) {
+ return nil, nil
+ })
+ }
+ }()
+
+ return deltaChan, wrappedErrChan
+}
+
+
+
package google
+
+import (
+ "encoding/json"
+ "fmt"
+ "math/rand"
+ "time"
+
+ "google.golang.org/genai"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+)
+
+// parseTools converts generic tool definitions from req.Tools (JSON) to Google's []*genai.Tool format.
+func parseTools(req *api.ResponseRequest) ([]*genai.Tool, error) {
+ if req.Tools == nil || len(req.Tools) == 0 {
+ return nil, nil
+ }
+
+ // Unmarshal to slice of tool definitions
+ var toolDefs []map[string]interface{}
+ if err := json.Unmarshal(req.Tools, &toolDefs); err != nil {
+ return nil, fmt.Errorf("unmarshal tools: %w", err)
+ }
+
+ var functionDeclarations []*genai.FunctionDeclaration
+
+ for _, toolDef := range toolDefs {
+ // Extract function details
+ // Support both flat format (name/description/parameters at top level)
+ // and nested format (under "function" key)
+ var name, description string
+ var parameters interface{}
+
+ if functionData, ok := toolDef["function"].(map[string]interface{}); ok {
+ // Nested format: {"type": "function", "function": {...}}
+ name, _ = functionData["name"].(string)
+ description, _ = functionData["description"].(string)
+ parameters = functionData["parameters"]
+ } else {
+ // Flat format: {"type": "function", "name": "...", ...}
+ name, _ = toolDef["name"].(string)
+ description, _ = toolDef["description"].(string)
+ parameters = toolDef["parameters"]
+ }
+
+ if name == "" {
+ continue
+ }
+
+ // Create function declaration
+ funcDecl := &genai.FunctionDeclaration{
+ Name: name,
+ Description: description,
+ }
+
+ // Google accepts parameters as raw JSON schema
+ if parameters != nil {
+ funcDecl.ParametersJsonSchema = parameters
+ }
+
+ functionDeclarations = append(functionDeclarations, funcDecl)
+ }
+
+ // Return single Tool with all function declarations
+ if len(functionDeclarations) > 0 {
+ return []*genai.Tool{{FunctionDeclarations: functionDeclarations}}, nil
+ }
+
+ return nil, nil
+}
+
+// parseToolChoice converts req.ToolChoice to Google's ToolConfig with FunctionCallingConfig.
+func parseToolChoice(req *api.ResponseRequest) (*genai.ToolConfig, error) {
+ if req.ToolChoice == nil || len(req.ToolChoice) == 0 {
+ return nil, nil
+ }
+
+ var choice interface{}
+ if err := json.Unmarshal(req.ToolChoice, &choice); err != nil {
+ return nil, fmt.Errorf("unmarshal tool_choice: %w", err)
+ }
+
+ config := &genai.ToolConfig{
+ FunctionCallingConfig: &genai.FunctionCallingConfig{},
+ }
+
+ // Handle string values: "auto", "none", "required"/"any"
+ if str, ok := choice.(string); ok {
+ switch str {
+ case "auto":
+ config.FunctionCallingConfig.Mode = genai.FunctionCallingConfigModeAuto
+ case "none":
+ config.FunctionCallingConfig.Mode = genai.FunctionCallingConfigModeNone
+ case "required", "any":
+ config.FunctionCallingConfig.Mode = genai.FunctionCallingConfigModeAny
+ default:
+ return nil, fmt.Errorf("unknown tool_choice string: %s", str)
+ }
+ return config, nil
+ }
+
+ // Handle object format: {"type": "function", "function": {"name": "..."}}
+ if obj, ok := choice.(map[string]interface{}); ok {
+ if typeVal, ok := obj["type"].(string); ok && typeVal == "function" {
+ config.FunctionCallingConfig.Mode = genai.FunctionCallingConfigModeAny
+ if funcObj, ok := obj["function"].(map[string]interface{}); ok {
+ if name, ok := funcObj["name"].(string); ok {
+ config.FunctionCallingConfig.AllowedFunctionNames = []string{name}
+ }
+ }
+ return config, nil
+ }
+ }
+
+ return nil, fmt.Errorf("unsupported tool_choice format")
+}
+
+// extractToolCalls extracts tool calls from Google's response format to generic api.ToolCall slice.
+func extractToolCalls(resp *genai.GenerateContentResponse) []api.ToolCall {
+ var toolCalls []api.ToolCall
+
+ for _, candidate := range resp.Candidates {
+ if candidate.Content == nil {
+ continue
+ }
+
+ for _, part := range candidate.Content.Parts {
+ if part == nil || part.FunctionCall == nil {
+ continue
+ }
+
+ // Extract function call details
+ fc := part.FunctionCall
+
+ // Marshal arguments to JSON string
+ var argsJSON string
+ if fc.Args != nil {
+ argsBytes, err := json.Marshal(fc.Args)
+ if err == nil {
+ argsJSON = string(argsBytes)
+ } else {
+ // Fallback to empty object
+ argsJSON = "{}"
+ }
+ } else {
+ argsJSON = "{}"
+ }
+
+ // Generate ID if Google doesn't provide one
+ callID := fc.ID
+ if callID == "" {
+ callID = fmt.Sprintf("call_%s", generateRandomID())
+ }
+
+ toolCalls = append(toolCalls, api.ToolCall{
+ ID: callID,
+ Name: fc.Name,
+ Arguments: argsJSON,
+ })
+ }
+ }
+
+ return toolCalls
+}
+
+// extractToolCallDelta extracts streaming tool call information from response parts.
+func extractToolCallDelta(part *genai.Part, index int) *api.ToolCallDelta {
+ if part == nil || part.FunctionCall == nil {
+ return nil
+ }
+
+ fc := part.FunctionCall
+
+ // Marshal arguments to JSON string
+ var argsJSON string
+ if fc.Args != nil {
+ argsBytes, err := json.Marshal(fc.Args)
+ if err == nil {
+ argsJSON = string(argsBytes)
+ } else {
+ argsJSON = "{}"
+ }
+ } else {
+ argsJSON = "{}"
+ }
+
+ // Generate ID if Google doesn't provide one
+ callID := fc.ID
+ if callID == "" {
+ callID = fmt.Sprintf("call_%s", generateRandomID())
+ }
+
+ return &api.ToolCallDelta{
+ Index: index,
+ ID: callID,
+ Name: fc.Name,
+ Arguments: argsJSON,
+ }
+}
+
+// generateRandomID generates a random alphanumeric ID
+func generateRandomID() string {
+ const charset = "abcdefghijklmnopqrstuvwxyz0123456789"
+ const length = 24
+ rng := rand.New(rand.NewSource(time.Now().UnixNano()))
+ b := make([]byte, length)
+ for i := range b {
+ b[i] = charset[rng.Intn(len(charset))]
+ }
+ return string(b)
+}
+
+
+
package google
+
+import (
+ "context"
+ "encoding/json"
+ "fmt"
+
+ "github.com/google/uuid"
+ "google.golang.org/genai"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/config"
+)
+
+const Name = "google"
+
+// Provider implements the Google Generative AI integration.
+type Provider struct {
+ cfg config.ProviderConfig
+ client *genai.Client
+}
+
+// New constructs a Provider using the Google AI API with API key authentication.
+func New(cfg config.ProviderConfig) (*Provider, error) {
+ var client *genai.Client
+ if cfg.APIKey != "" {
+ var err error
+ client, err = genai.NewClient(context.Background(), &genai.ClientConfig{
+ APIKey: cfg.APIKey,
+ })
+ if err != nil {
+ return nil, fmt.Errorf("failed to create google client: %w", err)
+ }
+ }
+ return &Provider{
+ cfg: cfg,
+ client: client,
+ }, nil
+}
+
+// NewVertexAI constructs a Provider targeting Vertex AI.
+// Vertex AI uses the same genai SDK but with GCP project/location configuration
+// and Application Default Credentials (ADC) or service account authentication.
+func NewVertexAI(vertexCfg config.VertexAIConfig) (*Provider, error) {
+ var client *genai.Client
+ if vertexCfg.Project != "" && vertexCfg.Location != "" {
+ var err error
+ client, err = genai.NewClient(context.Background(), &genai.ClientConfig{
+ Project: vertexCfg.Project,
+ Location: vertexCfg.Location,
+ Backend: genai.BackendVertexAI,
+ })
+ if err != nil {
+ return nil, fmt.Errorf("failed to create vertex ai client: %w", err)
+ }
+ }
+ return &Provider{
+ cfg: config.ProviderConfig{
+ // Vertex AI doesn't use API key, but set empty for consistency
+ APIKey: "",
+ },
+ client: client,
+ }, nil
+}
+
+func (p *Provider) Name() string { return Name }
+
+// Generate routes the request to Gemini and returns a ProviderResult.
+func (p *Provider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
+ if p.client == nil {
+ return nil, fmt.Errorf("google client not initialized")
+ }
+
+ model := req.Model
+
+ contents, systemText := convertMessages(messages)
+
+ // Parse tools if present
+ var tools []*genai.Tool
+ if req.Tools != nil && len(req.Tools) > 0 {
+ var err error
+ tools, err = parseTools(req)
+ if err != nil {
+ return nil, fmt.Errorf("parse tools: %w", err)
+ }
+ }
+
+ // Parse tool_choice if present
+ var toolConfig *genai.ToolConfig
+ if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
+ var err error
+ toolConfig, err = parseToolChoice(req)
+ if err != nil {
+ return nil, fmt.Errorf("parse tool_choice: %w", err)
+ }
+ }
+
+ config := buildConfig(systemText, req, tools, toolConfig)
+
+ resp, err := p.client.Models.GenerateContent(ctx, model, contents, config)
+ if err != nil {
+ return nil, fmt.Errorf("google api error: %w", err)
+ }
+
+ var text string
+ if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
+ for _, part := range resp.Candidates[0].Content.Parts {
+ if part != nil {
+ text += part.Text
+ }
+ }
+ }
+
+ var toolCalls []api.ToolCall
+ if len(resp.Candidates) > 0 {
+ toolCalls = extractToolCalls(resp)
+ }
+
+ var inputTokens, outputTokens int
+ if resp.UsageMetadata != nil {
+ inputTokens = int(resp.UsageMetadata.PromptTokenCount)
+ outputTokens = int(resp.UsageMetadata.CandidatesTokenCount)
+ }
+
+ return &api.ProviderResult{
+ ID: uuid.NewString(),
+ Model: model,
+ Text: text,
+ ToolCalls: toolCalls,
+ Usage: api.Usage{
+ InputTokens: inputTokens,
+ OutputTokens: outputTokens,
+ TotalTokens: inputTokens + outputTokens,
+ },
+ }, nil
+}
+
+// GenerateStream handles streaming requests to Google.
+func (p *Provider) GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error) {
+ deltaChan := make(chan *api.ProviderStreamDelta)
+ errChan := make(chan error, 1)
+
+ go func() {
+ defer close(deltaChan)
+ defer close(errChan)
+
+ if p.client == nil {
+ errChan <- fmt.Errorf("google client not initialized")
+ return
+ }
+
+ model := req.Model
+
+ contents, systemText := convertMessages(messages)
+
+ // Parse tools if present
+ var tools []*genai.Tool
+ if req.Tools != nil && len(req.Tools) > 0 {
+ var err error
+ tools, err = parseTools(req)
+ if err != nil {
+ errChan <- fmt.Errorf("parse tools: %w", err)
+ return
+ }
+ }
+
+ // Parse tool_choice if present
+ var toolConfig *genai.ToolConfig
+ if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
+ var err error
+ toolConfig, err = parseToolChoice(req)
+ if err != nil {
+ errChan <- fmt.Errorf("parse tool_choice: %w", err)
+ return
+ }
+ }
+
+ config := buildConfig(systemText, req, tools, toolConfig)
+
+ stream := p.client.Models.GenerateContentStream(ctx, model, contents, config)
+
+ for resp, err := range stream {
+ if err != nil {
+ errChan <- fmt.Errorf("google stream error: %w", err)
+ return
+ }
+
+ if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
+ for partIndex, part := range resp.Candidates[0].Content.Parts {
+ if part != nil {
+ // Handle text content
+ if part.Text != "" {
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{Text: part.Text}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ }
+
+ // Handle tool call content
+ if part.FunctionCall != nil {
+ delta := extractToolCallDelta(part, partIndex)
+ if delta != nil {
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{ToolCallDelta: delta}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ }
+ }
+ }
+ }
+ }
+ }
+
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{Done: true}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ }
+ }()
+
+ return deltaChan, errChan
+}
+
+// convertMessages splits messages into Gemini contents and system text.
+func convertMessages(messages []api.Message) ([]*genai.Content, string) {
+ var contents []*genai.Content
+ var systemText string
+
+ // Build a map of CallID -> Name from assistant tool calls
+ // This allows us to look up function names when processing tool results
+ callIDToName := make(map[string]string)
+ for _, msg := range messages {
+ if msg.Role == "assistant" || msg.Role == "model" {
+ for _, tc := range msg.ToolCalls {
+ if tc.ID != "" && tc.Name != "" {
+ callIDToName[tc.ID] = tc.Name
+ }
+ }
+ }
+ }
+
+ for _, msg := range messages {
+ if msg.Role == "system" || msg.Role == "developer" {
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ systemText += block.Text
+ }
+ }
+ continue
+ }
+
+ if msg.Role == "tool" {
+ // Tool results are sent as FunctionResponse in user role message
+ var output string
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ output += block.Text
+ }
+ }
+
+ // Parse output as JSON map, or wrap in {"output": "..."} if not JSON
+ var responseMap map[string]any
+ if err := json.Unmarshal([]byte(output), &responseMap); err != nil {
+ // Not JSON, wrap it
+ responseMap = map[string]any{"output": output}
+ }
+
+ // Get function name from message or look it up from CallID
+ name := msg.Name
+ if name == "" && msg.CallID != "" {
+ name = callIDToName[msg.CallID]
+ }
+
+ // Create FunctionResponse part with CallID and Name from message
+ part := &genai.Part{
+ FunctionResponse: &genai.FunctionResponse{
+ ID: msg.CallID,
+ Name: name, // Name is required by Google
+ Response: responseMap,
+ },
+ }
+
+ // Add to user role message
+ contents = append(contents, &genai.Content{
+ Role: "user",
+ Parts: []*genai.Part{part},
+ })
+ continue
+ }
+
+ var parts []*genai.Part
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ parts = append(parts, genai.NewPartFromText(block.Text))
+ }
+ }
+
+ // Add tool calls for assistant messages
+ if msg.Role == "assistant" || msg.Role == "model" {
+ for _, tc := range msg.ToolCalls {
+ // Parse arguments JSON into map
+ var args map[string]any
+ if err := json.Unmarshal([]byte(tc.Arguments), &args); err != nil {
+ // If unmarshal fails, skip this tool call
+ continue
+ }
+
+ // Create FunctionCall part
+ parts = append(parts, &genai.Part{
+ FunctionCall: &genai.FunctionCall{
+ ID: tc.ID,
+ Name: tc.Name,
+ Args: args,
+ },
+ })
+ }
+ }
+
+ role := "user"
+ if msg.Role == "assistant" || msg.Role == "model" {
+ role = "model"
+ }
+
+ contents = append(contents, &genai.Content{
+ Role: role,
+ Parts: parts,
+ })
+ }
+
+ return contents, systemText
+}
+
+// buildConfig constructs a GenerateContentConfig from system text and request params.
+func buildConfig(systemText string, req *api.ResponseRequest, tools []*genai.Tool, toolConfig *genai.ToolConfig) *genai.GenerateContentConfig {
+ var cfg *genai.GenerateContentConfig
+
+ needsCfg := systemText != "" || req.MaxOutputTokens != nil || req.Temperature != nil || req.TopP != nil || tools != nil || toolConfig != nil
+ if !needsCfg {
+ return nil
+ }
+
+ cfg = &genai.GenerateContentConfig{}
+
+ if systemText != "" {
+ cfg.SystemInstruction = &genai.Content{
+ Parts: []*genai.Part{genai.NewPartFromText(systemText)},
+ }
+ }
+
+ if req.MaxOutputTokens != nil {
+ cfg.MaxOutputTokens = int32(*req.MaxOutputTokens)
+ }
+
+ if req.Temperature != nil {
+ t := float32(*req.Temperature)
+ cfg.Temperature = &t
+ }
+
+ if req.TopP != nil {
+ tp := float32(*req.TopP)
+ cfg.TopP = &tp
+ }
+
+ if tools != nil {
+ cfg.Tools = tools
+ }
+
+ if toolConfig != nil {
+ cfg.ToolConfig = toolConfig
+ }
+
+ return cfg
+}
+
+func chooseModel(requested, defaultModel string) string {
+ if requested != "" {
+ return requested
+ }
+ if defaultModel != "" {
+ return defaultModel
+ }
+ return "gemini-2.0-flash-exp"
+}
+
+
+
package openai
+
+import (
+ "encoding/json"
+ "fmt"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/openai/openai-go/v3"
+ "github.com/openai/openai-go/v3/shared"
+)
+
+// parseTools converts Open Responses tools to OpenAI format
+func parseTools(req *api.ResponseRequest) ([]openai.ChatCompletionToolUnionParam, error) {
+ if req.Tools == nil || len(req.Tools) == 0 {
+ return nil, nil
+ }
+
+ var toolDefs []map[string]interface{}
+ if err := json.Unmarshal(req.Tools, &toolDefs); err != nil {
+ return nil, fmt.Errorf("unmarshal tools: %w", err)
+ }
+
+ var tools []openai.ChatCompletionToolUnionParam
+ for _, td := range toolDefs {
+ // Convert Open Responses tool to OpenAI ChatCompletionFunctionToolParam
+ // Extract: name, description, parameters
+ name, _ := td["name"].(string)
+ desc, _ := td["description"].(string)
+ params, _ := td["parameters"].(map[string]interface{})
+
+ funcDef := shared.FunctionDefinitionParam{
+ Name: name,
+ }
+
+ if desc != "" {
+ funcDef.Description = openai.String(desc)
+ }
+
+ if params != nil {
+ funcDef.Parameters = shared.FunctionParameters(params)
+ }
+
+ tools = append(tools, openai.ChatCompletionFunctionTool(funcDef))
+ }
+
+ return tools, nil
+}
+
+// parseToolChoice converts Open Responses tool_choice to OpenAI format
+func parseToolChoice(req *api.ResponseRequest) (openai.ChatCompletionToolChoiceOptionUnionParam, error) {
+ var result openai.ChatCompletionToolChoiceOptionUnionParam
+
+ if req.ToolChoice == nil || len(req.ToolChoice) == 0 {
+ return result, nil
+ }
+
+ var choice interface{}
+ if err := json.Unmarshal(req.ToolChoice, &choice); err != nil {
+ return result, fmt.Errorf("unmarshal tool_choice: %w", err)
+ }
+
+ // Handle string values: "auto", "none", "required"
+ if str, ok := choice.(string); ok {
+ result.OfAuto = openai.String(str)
+ return result, nil
+ }
+
+ // Handle specific function selection: {"type": "function", "function": {"name": "..."}}
+ if obj, ok := choice.(map[string]interface{}); ok {
+ funcObj, _ := obj["function"].(map[string]interface{})
+ name, _ := funcObj["name"].(string)
+
+ return openai.ToolChoiceOptionFunctionToolChoice(
+ openai.ChatCompletionNamedToolChoiceFunctionParam{
+ Name: name,
+ },
+ ), nil
+ }
+
+ return result, fmt.Errorf("invalid tool_choice format")
+}
+
+// extractToolCalls converts OpenAI tool calls to api.ToolCall
+func extractToolCalls(message openai.ChatCompletionMessage) []api.ToolCall {
+ if len(message.ToolCalls) == 0 {
+ return nil
+ }
+
+ var toolCalls []api.ToolCall
+ for _, tc := range message.ToolCalls {
+ toolCalls = append(toolCalls, api.ToolCall{
+ ID: tc.ID,
+ Name: tc.Function.Name,
+ Arguments: tc.Function.Arguments,
+ })
+ }
+ return toolCalls
+}
+
+// extractToolCallDelta extracts tool call delta from streaming chunk choice
+func extractToolCallDelta(choice openai.ChatCompletionChunkChoice) *api.ToolCallDelta {
+ if len(choice.Delta.ToolCalls) == 0 {
+ return nil
+ }
+
+ // OpenAI sends tool calls with index in the delta
+ for _, tc := range choice.Delta.ToolCalls {
+ return &api.ToolCallDelta{
+ Index: int(tc.Index),
+ ID: tc.ID,
+ Name: tc.Function.Name,
+ Arguments: tc.Function.Arguments,
+ }
+ }
+
+ return nil
+}
+
+
+
package openai
+
+import (
+ "context"
+ "fmt"
+
+ "github.com/openai/openai-go/v3"
+ "github.com/openai/openai-go/v3/azure"
+ "github.com/openai/openai-go/v3/option"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/config"
+)
+
+const Name = "openai"
+
+// Provider implements the OpenAI SDK integration.
+// It supports both direct OpenAI API and Azure-hosted endpoints.
+type Provider struct {
+ cfg config.ProviderConfig
+ client *openai.Client
+ azure bool
+}
+
+// New constructs a Provider for the direct OpenAI API.
+func New(cfg config.ProviderConfig) *Provider {
+ var client *openai.Client
+ if cfg.APIKey != "" {
+ c := openai.NewClient(option.WithAPIKey(cfg.APIKey))
+ client = &c
+ }
+ return &Provider{
+ cfg: cfg,
+ client: client,
+ }
+}
+
+// NewAzure constructs a Provider targeting Azure OpenAI.
+// Azure OpenAI uses the OpenAI SDK with the azure subpackage for proper
+// endpoint routing, api-version query parameter, and API key header.
+func NewAzure(azureCfg config.AzureOpenAIConfig) *Provider {
+ var client *openai.Client
+ if azureCfg.APIKey != "" && azureCfg.Endpoint != "" {
+ apiVersion := azureCfg.APIVersion
+ if apiVersion == "" {
+ apiVersion = "2024-12-01-preview"
+ }
+ c := openai.NewClient(
+ azure.WithEndpoint(azureCfg.Endpoint, apiVersion),
+ azure.WithAPIKey(azureCfg.APIKey),
+ )
+ client = &c
+ }
+ return &Provider{
+ cfg: config.ProviderConfig{
+ APIKey: azureCfg.APIKey,
+ },
+ client: client,
+ azure: true,
+ }
+}
+
+// Name returns the provider identifier.
+func (p *Provider) Name() string { return Name }
+
+// Generate routes the request to OpenAI and returns a ProviderResult.
+func (p *Provider) Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error) {
+ if p.cfg.APIKey == "" {
+ return nil, fmt.Errorf("openai api key missing")
+ }
+ if p.client == nil {
+ return nil, fmt.Errorf("openai client not initialized")
+ }
+
+ // Convert messages to OpenAI format
+ oaiMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages))
+ for _, msg := range messages {
+ var content string
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ content += block.Text
+ }
+ }
+
+ switch msg.Role {
+ case "user":
+ oaiMessages = append(oaiMessages, openai.UserMessage(content))
+ case "assistant":
+ // If assistant message has tool calls, include them
+ if len(msg.ToolCalls) > 0 {
+ toolCalls := make([]openai.ChatCompletionMessageToolCallUnionParam, len(msg.ToolCalls))
+ for i, tc := range msg.ToolCalls {
+ toolCalls[i] = openai.ChatCompletionMessageToolCallUnionParam{
+ OfFunction: &openai.ChatCompletionMessageFunctionToolCallParam{
+ ID: tc.ID,
+ Function: openai.ChatCompletionMessageFunctionToolCallFunctionParam{
+ Name: tc.Name,
+ Arguments: tc.Arguments,
+ },
+ },
+ }
+ }
+ msgParam := openai.ChatCompletionAssistantMessageParam{
+ ToolCalls: toolCalls,
+ }
+ if content != "" {
+ msgParam.Content.OfString = openai.String(content)
+ }
+ oaiMessages = append(oaiMessages, openai.ChatCompletionMessageParamUnion{
+ OfAssistant: &msgParam,
+ })
+ } else {
+ oaiMessages = append(oaiMessages, openai.AssistantMessage(content))
+ }
+ case "system":
+ oaiMessages = append(oaiMessages, openai.SystemMessage(content))
+ case "developer":
+ oaiMessages = append(oaiMessages, openai.SystemMessage(content))
+ case "tool":
+ oaiMessages = append(oaiMessages, openai.ToolMessage(content, msg.CallID))
+ }
+ }
+
+ params := openai.ChatCompletionNewParams{
+ Model: openai.ChatModel(req.Model),
+ Messages: oaiMessages,
+ }
+ if req.MaxOutputTokens != nil {
+ params.MaxTokens = openai.Int(int64(*req.MaxOutputTokens))
+ }
+ if req.Temperature != nil {
+ params.Temperature = openai.Float(*req.Temperature)
+ }
+ if req.TopP != nil {
+ params.TopP = openai.Float(*req.TopP)
+ }
+
+ // Add tools if present
+ if req.Tools != nil && len(req.Tools) > 0 {
+ tools, err := parseTools(req)
+ if err != nil {
+ return nil, fmt.Errorf("parse tools: %w", err)
+ }
+ params.Tools = tools
+ }
+
+ // Add tool_choice if present
+ if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
+ toolChoice, err := parseToolChoice(req)
+ if err != nil {
+ return nil, fmt.Errorf("parse tool_choice: %w", err)
+ }
+ params.ToolChoice = toolChoice
+ }
+
+ // Add parallel_tool_calls if specified
+ if req.ParallelToolCalls != nil {
+ params.ParallelToolCalls = openai.Bool(*req.ParallelToolCalls)
+ }
+
+ // Call OpenAI API
+ resp, err := p.client.Chat.Completions.New(ctx, params)
+ if err != nil {
+ return nil, fmt.Errorf("openai api error: %w", err)
+ }
+
+ var combinedText string
+ var toolCalls []api.ToolCall
+
+ for _, choice := range resp.Choices {
+ combinedText += choice.Message.Content
+ if len(choice.Message.ToolCalls) > 0 {
+ toolCalls = append(toolCalls, extractToolCalls(choice.Message)...)
+ }
+ }
+
+ return &api.ProviderResult{
+ ID: resp.ID,
+ Model: resp.Model,
+ Text: combinedText,
+ ToolCalls: toolCalls,
+ Usage: api.Usage{
+ InputTokens: int(resp.Usage.PromptTokens),
+ OutputTokens: int(resp.Usage.CompletionTokens),
+ TotalTokens: int(resp.Usage.TotalTokens),
+ },
+ }, nil
+}
+
+// GenerateStream handles streaming requests to OpenAI.
+func (p *Provider) GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error) {
+ deltaChan := make(chan *api.ProviderStreamDelta)
+ errChan := make(chan error, 1)
+
+ go func() {
+ defer close(deltaChan)
+ defer close(errChan)
+
+ if p.cfg.APIKey == "" {
+ errChan <- fmt.Errorf("openai api key missing")
+ return
+ }
+ if p.client == nil {
+ errChan <- fmt.Errorf("openai client not initialized")
+ return
+ }
+
+ // Convert messages to OpenAI format
+ oaiMessages := make([]openai.ChatCompletionMessageParamUnion, 0, len(messages))
+ for _, msg := range messages {
+ var content string
+ for _, block := range msg.Content {
+ if block.Type == "input_text" || block.Type == "output_text" {
+ content += block.Text
+ }
+ }
+
+ switch msg.Role {
+ case "user":
+ oaiMessages = append(oaiMessages, openai.UserMessage(content))
+ case "assistant":
+ // If assistant message has tool calls, include them
+ if len(msg.ToolCalls) > 0 {
+ toolCalls := make([]openai.ChatCompletionMessageToolCallUnionParam, len(msg.ToolCalls))
+ for i, tc := range msg.ToolCalls {
+ toolCalls[i] = openai.ChatCompletionMessageToolCallUnionParam{
+ OfFunction: &openai.ChatCompletionMessageFunctionToolCallParam{
+ ID: tc.ID,
+ Function: openai.ChatCompletionMessageFunctionToolCallFunctionParam{
+ Name: tc.Name,
+ Arguments: tc.Arguments,
+ },
+ },
+ }
+ }
+ msgParam := openai.ChatCompletionAssistantMessageParam{
+ ToolCalls: toolCalls,
+ }
+ if content != "" {
+ msgParam.Content.OfString = openai.String(content)
+ }
+ oaiMessages = append(oaiMessages, openai.ChatCompletionMessageParamUnion{
+ OfAssistant: &msgParam,
+ })
+ } else {
+ oaiMessages = append(oaiMessages, openai.AssistantMessage(content))
+ }
+ case "system":
+ oaiMessages = append(oaiMessages, openai.SystemMessage(content))
+ case "developer":
+ oaiMessages = append(oaiMessages, openai.SystemMessage(content))
+ case "tool":
+ oaiMessages = append(oaiMessages, openai.ToolMessage(content, msg.CallID))
+ }
+ }
+
+ params := openai.ChatCompletionNewParams{
+ Model: openai.ChatModel(req.Model),
+ Messages: oaiMessages,
+ }
+ if req.MaxOutputTokens != nil {
+ params.MaxTokens = openai.Int(int64(*req.MaxOutputTokens))
+ }
+ if req.Temperature != nil {
+ params.Temperature = openai.Float(*req.Temperature)
+ }
+ if req.TopP != nil {
+ params.TopP = openai.Float(*req.TopP)
+ }
+
+ // Add tools if present
+ if req.Tools != nil && len(req.Tools) > 0 {
+ tools, err := parseTools(req)
+ if err != nil {
+ errChan <- fmt.Errorf("parse tools: %w", err)
+ return
+ }
+ params.Tools = tools
+ }
+
+ // Add tool_choice if present
+ if req.ToolChoice != nil && len(req.ToolChoice) > 0 {
+ toolChoice, err := parseToolChoice(req)
+ if err != nil {
+ errChan <- fmt.Errorf("parse tool_choice: %w", err)
+ return
+ }
+ params.ToolChoice = toolChoice
+ }
+
+ // Add parallel_tool_calls if specified
+ if req.ParallelToolCalls != nil {
+ params.ParallelToolCalls = openai.Bool(*req.ParallelToolCalls)
+ }
+
+ // Create streaming request
+ stream := p.client.Chat.Completions.NewStreaming(ctx, params)
+
+ // Process stream
+ for stream.Next() {
+ chunk := stream.Current()
+
+ for _, choice := range chunk.Choices {
+ // Handle text content
+ if choice.Delta.Content != "" {
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{
+ ID: chunk.ID,
+ Model: chunk.Model,
+ Text: choice.Delta.Content,
+ }:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ }
+
+ // Handle tool call deltas
+ if len(choice.Delta.ToolCalls) > 0 {
+ delta := extractToolCallDelta(choice)
+ if delta != nil {
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{
+ ID: chunk.ID,
+ Model: chunk.Model,
+ ToolCallDelta: delta,
+ }:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ return
+ }
+ }
+ }
+ }
+ }
+
+ if err := stream.Err(); err != nil {
+ errChan <- fmt.Errorf("openai stream error: %w", err)
+ return
+ }
+
+ // Send final delta
+ select {
+ case deltaChan <- &api.ProviderStreamDelta{Done: true}:
+ case <-ctx.Done():
+ errChan <- ctx.Err()
+ }
+ }()
+
+ return deltaChan, errChan
+}
+
+func chooseModel(requested, defaultModel string) string {
+ if requested != "" {
+ return requested
+ }
+ if defaultModel != "" {
+ return defaultModel
+ }
+ return "gpt-4o-mini"
+}
+
+
+
package providers
+
+import (
+ "context"
+ "fmt"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/config"
+ anthropicprovider "github.com/ajac-zero/latticelm/internal/providers/anthropic"
+ googleprovider "github.com/ajac-zero/latticelm/internal/providers/google"
+ openaiprovider "github.com/ajac-zero/latticelm/internal/providers/openai"
+)
+
+// Provider represents a unified interface that each LLM provider must implement.
+type Provider interface {
+ Name() string
+ Generate(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (*api.ProviderResult, error)
+ GenerateStream(ctx context.Context, messages []api.Message, req *api.ResponseRequest) (<-chan *api.ProviderStreamDelta, <-chan error)
+}
+
+// Registry keeps track of registered providers and model-to-provider mappings.
+type Registry struct {
+ providers map[string]Provider
+ models map[string]string // model name -> provider entry name
+ providerModelIDs map[string]string // model name -> provider model ID
+ modelList []config.ModelEntry
+}
+
+// NewRegistry constructs provider implementations from configuration.
+func NewRegistry(entries map[string]config.ProviderEntry, models []config.ModelEntry) (*Registry, error) {
+ return NewRegistryWithCircuitBreaker(entries, models, nil)
+}
+
+// NewRegistryWithCircuitBreaker constructs provider implementations with circuit breaker support.
+// The onStateChange callback is invoked when circuit breaker state changes.
+func NewRegistryWithCircuitBreaker(
+ entries map[string]config.ProviderEntry,
+ models []config.ModelEntry,
+ onStateChange func(provider, from, to string),
+) (*Registry, error) {
+ reg := &Registry{
+ providers: make(map[string]Provider),
+ models: make(map[string]string),
+ providerModelIDs: make(map[string]string),
+ modelList: models,
+ }
+
+ // Use default circuit breaker configuration
+ cbConfig := DefaultCircuitBreakerConfig()
+ cbConfig.OnStateChange = onStateChange
+
+ for name, entry := range entries {
+ p, err := buildProvider(entry)
+ if err != nil {
+ return nil, fmt.Errorf("provider %q: %w", name, err)
+ }
+ if p != nil {
+ // Wrap provider with circuit breaker
+ reg.providers[name] = NewCircuitBreakerProvider(p, cbConfig)
+ }
+ }
+
+ for _, m := range models {
+ reg.models[m.Name] = m.Provider
+ if m.ProviderModelID != "" {
+ reg.providerModelIDs[m.Name] = m.ProviderModelID
+ }
+ }
+
+ if len(reg.providers) == 0 {
+ return nil, fmt.Errorf("no providers configured")
+ }
+
+ return reg, nil
+}
+
+func buildProvider(entry config.ProviderEntry) (Provider, error) {
+ // Vertex AI doesn't require APIKey, so check for it separately
+ if entry.Type != "vertexai" && entry.APIKey == "" {
+ return nil, nil
+ }
+
+ switch entry.Type {
+ case "openai":
+ return openaiprovider.New(config.ProviderConfig{
+ APIKey: entry.APIKey,
+ Endpoint: entry.Endpoint,
+ }), nil
+ case "azureopenai":
+ if entry.Endpoint == "" {
+ return nil, fmt.Errorf("endpoint is required for azureopenai")
+ }
+ return openaiprovider.NewAzure(config.AzureOpenAIConfig{
+ APIKey: entry.APIKey,
+ Endpoint: entry.Endpoint,
+ APIVersion: entry.APIVersion,
+ }), nil
+ case "anthropic":
+ return anthropicprovider.New(config.ProviderConfig{
+ APIKey: entry.APIKey,
+ Endpoint: entry.Endpoint,
+ }), nil
+ case "azureanthropic":
+ if entry.Endpoint == "" {
+ return nil, fmt.Errorf("endpoint is required for azureanthropic")
+ }
+ return anthropicprovider.NewAzure(config.AzureAnthropicConfig{
+ APIKey: entry.APIKey,
+ Endpoint: entry.Endpoint,
+ }), nil
+ case "google":
+ return googleprovider.New(config.ProviderConfig{
+ APIKey: entry.APIKey,
+ Endpoint: entry.Endpoint,
+ })
+ case "vertexai":
+ if entry.Project == "" || entry.Location == "" {
+ return nil, fmt.Errorf("project and location are required for vertexai")
+ }
+ return googleprovider.NewVertexAI(config.VertexAIConfig{
+ Project: entry.Project,
+ Location: entry.Location,
+ })
+ default:
+ return nil, fmt.Errorf("unknown provider type %q", entry.Type)
+ }
+}
+
+// Get returns provider by entry name.
+func (r *Registry) Get(name string) (Provider, bool) {
+ p, ok := r.providers[name]
+ return p, ok
+}
+
+// Models returns the list of configured models and their provider entry names.
+func (r *Registry) Models() []struct{ Provider, Model string } {
+ var out []struct{ Provider, Model string }
+ for _, m := range r.modelList {
+ out = append(out, struct{ Provider, Model string }{Provider: m.Provider, Model: m.Name})
+ }
+ return out
+}
+
+// ResolveModelID returns the provider_model_id for a model, falling back to the model name itself.
+func (r *Registry) ResolveModelID(model string) string {
+ if id, ok := r.providerModelIDs[model]; ok {
+ return id
+ }
+ return model
+}
+
+// Default returns the provider for the given model name.
+func (r *Registry) Default(model string) (Provider, error) {
+ if model != "" {
+ if providerName, ok := r.models[model]; ok {
+ if p, ok := r.providers[providerName]; ok {
+ return p, nil
+ }
+ }
+ }
+
+ for _, p := range r.providers {
+ return p, nil
+ }
+
+ return nil, fmt.Errorf("no providers available")
+}
+
+
+
package ratelimit
+
+import (
+ "log/slog"
+ "net/http"
+ "sync"
+ "time"
+
+ "golang.org/x/time/rate"
+)
+
+// Config defines rate limiting configuration.
+type Config struct {
+ // RequestsPerSecond is the number of requests allowed per second per IP.
+ RequestsPerSecond float64
+ // Burst is the maximum burst size allowed.
+ Burst int
+ // Enabled controls whether rate limiting is active.
+ Enabled bool
+}
+
+// Middleware provides per-IP rate limiting using token bucket algorithm.
+type Middleware struct {
+ limiters map[string]*rate.Limiter
+ mu sync.RWMutex
+ config Config
+ logger *slog.Logger
+}
+
+// New creates a new rate limiting middleware.
+func New(config Config, logger *slog.Logger) *Middleware {
+ m := &Middleware{
+ limiters: make(map[string]*rate.Limiter),
+ config: config,
+ logger: logger,
+ }
+
+ // Start cleanup goroutine to remove old limiters
+ if config.Enabled {
+ go m.cleanupLimiters()
+ }
+
+ return m
+}
+
+// Handler wraps an http.Handler with rate limiting.
+func (m *Middleware) Handler(next http.Handler) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ if !m.config.Enabled {
+ next.ServeHTTP(w, r)
+ return
+ }
+
+ // Extract client IP (handle X-Forwarded-For for proxies)
+ ip := m.getClientIP(r)
+
+ limiter := m.getLimiter(ip)
+ if !limiter.Allow() {
+ m.logger.Warn("rate limit exceeded",
+ slog.String("ip", ip),
+ slog.String("path", r.URL.Path),
+ )
+ w.Header().Set("Content-Type", "application/json")
+ w.Header().Set("Retry-After", "1")
+ w.WriteHeader(http.StatusTooManyRequests)
+ w.Write([]byte(`{"error":"rate limit exceeded","message":"too many requests"}`))
+ return
+ }
+
+ next.ServeHTTP(w, r)
+ })
+}
+
+// getLimiter returns the rate limiter for a given IP, creating one if needed.
+func (m *Middleware) getLimiter(ip string) *rate.Limiter {
+ m.mu.RLock()
+ limiter, exists := m.limiters[ip]
+ m.mu.RUnlock()
+
+ if exists {
+ return limiter
+ }
+
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ // Double-check after acquiring write lock
+ limiter, exists = m.limiters[ip]
+ if exists {
+ return limiter
+ }
+
+ limiter = rate.NewLimiter(rate.Limit(m.config.RequestsPerSecond), m.config.Burst)
+ m.limiters[ip] = limiter
+ return limiter
+}
+
+// getClientIP extracts the client IP from the request.
+func (m *Middleware) getClientIP(r *http.Request) string {
+ // Check X-Forwarded-For header (for proxies/load balancers)
+ xff := r.Header.Get("X-Forwarded-For")
+ if xff != "" {
+ // X-Forwarded-For can be a comma-separated list, use the first IP
+ for idx := 0; idx < len(xff); idx++ {
+ if xff[idx] == ',' {
+ return xff[:idx]
+ }
+ }
+ return xff
+ }
+
+ // Check X-Real-IP header
+ if xri := r.Header.Get("X-Real-IP"); xri != "" {
+ return xri
+ }
+
+ // Fall back to RemoteAddr
+ return r.RemoteAddr
+}
+
+// cleanupLimiters periodically removes unused limiters to prevent memory leaks.
+func (m *Middleware) cleanupLimiters() {
+ ticker := time.NewTicker(5 * time.Minute)
+ defer ticker.Stop()
+
+ for range ticker.C {
+ m.mu.Lock()
+ // Clear all limiters periodically
+ // In production, you might want a more sophisticated LRU cache
+ m.limiters = make(map[string]*rate.Limiter)
+ m.mu.Unlock()
+
+ m.logger.Debug("cleaned up rate limiters")
+ }
+}
+
+
+
package server
+
+import (
+ "context"
+ "encoding/json"
+ "net/http"
+ "time"
+)
+
+// HealthStatus represents the health check response.
+type HealthStatus struct {
+ Status string `json:"status"`
+ Timestamp int64 `json:"timestamp"`
+ Checks map[string]string `json:"checks,omitempty"`
+}
+
+// handleHealth returns a basic health check endpoint.
+// This is suitable for Kubernetes liveness probes.
+func (s *GatewayServer) handleHealth(w http.ResponseWriter, r *http.Request) {
+ if r.Method != http.MethodGet {
+ http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+ return
+ }
+
+ status := HealthStatus{
+ Status: "healthy",
+ Timestamp: time.Now().Unix(),
+ }
+
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusOK)
+ if err := json.NewEncoder(w).Encode(status); err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to encode health response", "error", err.Error())
+ }
+}
+
+// handleReady returns a readiness check that verifies dependencies.
+// This is suitable for Kubernetes readiness probes and load balancer health checks.
+func (s *GatewayServer) handleReady(w http.ResponseWriter, r *http.Request) {
+ if r.Method != http.MethodGet {
+ http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+ return
+ }
+
+ checks := make(map[string]string)
+ allHealthy := true
+
+ // Check conversation store connectivity
+ ctx, cancel := context.WithTimeout(r.Context(), 2*time.Second)
+ defer cancel()
+
+ // Test conversation store by attempting a simple operation
+ testID := "health_check_test"
+ _, err := s.convs.Get(ctx, testID)
+ if err != nil {
+ checks["conversation_store"] = "unhealthy: " + err.Error()
+ allHealthy = false
+ } else {
+ checks["conversation_store"] = "healthy"
+ }
+
+ // Check if at least one provider is configured
+ models := s.registry.Models()
+ if len(models) == 0 {
+ checks["providers"] = "unhealthy: no providers configured"
+ allHealthy = false
+ } else {
+ checks["providers"] = "healthy"
+ }
+
+ _ = ctx // Use context if needed
+
+ status := HealthStatus{
+ Timestamp: time.Now().Unix(),
+ Checks: checks,
+ }
+
+ if allHealthy {
+ status.Status = "ready"
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusOK)
+ } else {
+ status.Status = "not_ready"
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusServiceUnavailable)
+ }
+
+ if err := json.NewEncoder(w).Encode(status); err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to encode ready response", "error", err.Error())
+ }
+}
+
+
+
package server
+
+import (
+ "fmt"
+ "log/slog"
+ "net/http"
+ "runtime/debug"
+
+ "github.com/ajac-zero/latticelm/internal/logger"
+)
+
+// MaxRequestBodyBytes is the maximum size allowed for request bodies (10MB)
+const MaxRequestBodyBytes = 10 * 1024 * 1024
+
+// PanicRecoveryMiddleware recovers from panics in HTTP handlers and logs them
+// instead of crashing the server. Returns 500 Internal Server Error to the client.
+func PanicRecoveryMiddleware(next http.Handler, log *slog.Logger) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ defer func() {
+ if err := recover(); err != nil {
+ // Capture stack trace
+ stack := debug.Stack()
+
+ // Log the panic with full context
+ log.ErrorContext(r.Context(), "panic recovered in HTTP handler",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("method", r.Method),
+ slog.String("path", r.URL.Path),
+ slog.String("remote_addr", r.RemoteAddr),
+ slog.Any("panic", err),
+ slog.String("stack", string(stack)),
+ )...,
+ )
+
+ // Return 500 to client
+ http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+ }
+ }()
+
+ next.ServeHTTP(w, r)
+ })
+}
+
+// RequestSizeLimitMiddleware enforces a maximum request body size to prevent
+// DoS attacks via oversized payloads. Requests exceeding the limit receive 413.
+func RequestSizeLimitMiddleware(next http.Handler, maxBytes int64) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ // Only limit body size for requests that have a body
+ if r.Method == http.MethodPost || r.Method == http.MethodPut || r.Method == http.MethodPatch {
+ // Wrap the request body with a size limiter
+ r.Body = http.MaxBytesReader(w, r.Body, maxBytes)
+ }
+
+ next.ServeHTTP(w, r)
+ })
+}
+
+// ErrorRecoveryMiddleware catches errors from MaxBytesReader and converts them
+// to proper HTTP error responses. This should be placed after RequestSizeLimitMiddleware
+// in the middleware chain.
+func ErrorRecoveryMiddleware(next http.Handler, log *slog.Logger) http.Handler {
+ return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ next.ServeHTTP(w, r)
+
+ // Check if the request body exceeded the size limit
+ // MaxBytesReader sets an error that we can detect on the next read attempt
+ // But we need to handle the error when it actually occurs during JSON decoding
+ // The JSON decoder will return the error, so we don't need special handling here
+ // This middleware is more for future extensibility
+ })
+}
+
+// WriteJSONError is a helper function to safely write JSON error responses,
+// handling any encoding errors that might occur.
+func WriteJSONError(w http.ResponseWriter, log *slog.Logger, message string, statusCode int) {
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(statusCode)
+
+ // Use fmt.Fprintf to write the error response
+ // This is safer than json.Encoder as we control the format
+ _, err := fmt.Fprintf(w, `{"error":{"message":"%s"}}`, message)
+ if err != nil {
+ // If we can't even write the error response, log it
+ log.Error("failed to write error response",
+ slog.String("original_message", message),
+ slog.Int("status_code", statusCode),
+ slog.String("write_error", err.Error()),
+ )
+ }
+}
+
+
+
package server
+
+import (
+ "encoding/json"
+ "errors"
+ "fmt"
+ "log/slog"
+ "net/http"
+ "strings"
+ "time"
+
+ "github.com/google/uuid"
+ "github.com/sony/gobreaker"
+
+ "github.com/ajac-zero/latticelm/internal/api"
+ "github.com/ajac-zero/latticelm/internal/conversation"
+ "github.com/ajac-zero/latticelm/internal/logger"
+ "github.com/ajac-zero/latticelm/internal/providers"
+)
+
+// ProviderRegistry is an interface for provider registries.
+type ProviderRegistry interface {
+ Get(name string) (providers.Provider, bool)
+ Models() []struct{ Provider, Model string }
+ ResolveModelID(model string) string
+ Default(model string) (providers.Provider, error)
+}
+
+// GatewayServer hosts the Open Responses API for the gateway.
+type GatewayServer struct {
+ registry ProviderRegistry
+ convs conversation.Store
+ logger *slog.Logger
+}
+
+// New creates a GatewayServer bound to the provider registry.
+func New(registry ProviderRegistry, convs conversation.Store, logger *slog.Logger) *GatewayServer {
+ return &GatewayServer{
+ registry: registry,
+ convs: convs,
+ logger: logger,
+ }
+}
+
+// isCircuitBreakerError checks if the error is from a circuit breaker.
+func isCircuitBreakerError(err error) bool {
+ return errors.Is(err, gobreaker.ErrOpenState) || errors.Is(err, gobreaker.ErrTooManyRequests)
+}
+
+// RegisterRoutes wires the HTTP handlers onto the provided mux.
+func (s *GatewayServer) RegisterRoutes(mux *http.ServeMux) {
+ mux.HandleFunc("/v1/responses", s.handleResponses)
+ mux.HandleFunc("/v1/models", s.handleModels)
+ mux.HandleFunc("/health", s.handleHealth)
+ mux.HandleFunc("/ready", s.handleReady)
+}
+
+func (s *GatewayServer) handleModels(w http.ResponseWriter, r *http.Request) {
+ if r.Method != http.MethodGet {
+ http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+ return
+ }
+
+ models := s.registry.Models()
+ var data []api.ModelInfo
+ for _, m := range models {
+ data = append(data, api.ModelInfo{
+ ID: m.Model,
+ Provider: m.Provider,
+ })
+ }
+
+ resp := api.ModelsResponse{
+ Object: "list",
+ Data: data,
+ }
+
+ w.Header().Set("Content-Type", "application/json")
+ if err := json.NewEncoder(w).Encode(resp); err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to encode models response",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("error", err.Error()),
+ )...,
+ )
+ }
+}
+
+func (s *GatewayServer) handleResponses(w http.ResponseWriter, r *http.Request) {
+ if r.Method != http.MethodPost {
+ http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
+ return
+ }
+
+ var req api.ResponseRequest
+ if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+ // Check if error is due to request size limit
+ if err.Error() == "http: request body too large" {
+ http.Error(w, "request body too large", http.StatusRequestEntityTooLarge)
+ return
+ }
+ http.Error(w, "invalid JSON payload", http.StatusBadRequest)
+ return
+ }
+
+ if err := req.Validate(); err != nil {
+ http.Error(w, err.Error(), http.StatusBadRequest)
+ return
+ }
+
+ // Normalize input to internal messages
+ inputMsgs := req.NormalizeInput()
+
+ // Build full message history from previous conversation
+ var historyMsgs []api.Message
+ if req.PreviousResponseID != nil && *req.PreviousResponseID != "" {
+ conv, err := s.convs.Get(r.Context(), *req.PreviousResponseID)
+ if err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to retrieve conversation",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("conversation_id", *req.PreviousResponseID),
+ slog.String("error", err.Error()),
+ )...,
+ )
+ http.Error(w, "error retrieving conversation", http.StatusInternalServerError)
+ return
+ }
+ if conv == nil {
+ s.logger.WarnContext(r.Context(), "conversation not found",
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("conversation_id", *req.PreviousResponseID),
+ )
+ http.Error(w, "conversation not found", http.StatusNotFound)
+ return
+ }
+ historyMsgs = conv.Messages
+ }
+
+ // Combined messages for conversation storage (history + new input, no instructions)
+ storeMsgs := make([]api.Message, 0, len(historyMsgs)+len(inputMsgs))
+ storeMsgs = append(storeMsgs, historyMsgs...)
+ storeMsgs = append(storeMsgs, inputMsgs...)
+
+ // Build provider messages: instructions + history + input
+ var providerMsgs []api.Message
+ if req.Instructions != nil && *req.Instructions != "" {
+ providerMsgs = append(providerMsgs, api.Message{
+ Role: "developer",
+ Content: []api.ContentBlock{{Type: "input_text", Text: *req.Instructions}},
+ })
+ }
+ providerMsgs = append(providerMsgs, storeMsgs...)
+
+ provider, err := s.resolveProvider(&req)
+ if err != nil {
+ http.Error(w, err.Error(), http.StatusBadGateway)
+ return
+ }
+
+ // Resolve provider_model_id (e.g., Azure deployment name)
+ resolvedReq := req
+ resolvedReq.Model = s.registry.ResolveModelID(req.Model)
+
+ if req.Stream {
+ s.handleStreamingResponse(w, r, provider, providerMsgs, &resolvedReq, &req, storeMsgs)
+ } else {
+ s.handleSyncResponse(w, r, provider, providerMsgs, &resolvedReq, &req, storeMsgs)
+ }
+}
+
+func (s *GatewayServer) handleSyncResponse(w http.ResponseWriter, r *http.Request, provider providers.Provider, providerMsgs []api.Message, resolvedReq *api.ResponseRequest, origReq *api.ResponseRequest, storeMsgs []api.Message) {
+ result, err := provider.Generate(r.Context(), providerMsgs, resolvedReq)
+ if err != nil {
+ s.logger.ErrorContext(r.Context(), "provider generation failed",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("provider", provider.Name()),
+ slog.String("model", resolvedReq.Model),
+ slog.String("error", err.Error()),
+ )...,
+ )
+
+ // Check if error is from circuit breaker
+ if isCircuitBreakerError(err) {
+ http.Error(w, "service temporarily unavailable - circuit breaker open", http.StatusServiceUnavailable)
+ } else {
+ http.Error(w, "provider error", http.StatusBadGateway)
+ }
+ return
+ }
+
+ responseID := generateID("resp_")
+
+ // Build assistant message for conversation store
+ assistantMsg := api.Message{
+ Role: "assistant",
+ Content: []api.ContentBlock{{Type: "output_text", Text: result.Text}},
+ ToolCalls: result.ToolCalls,
+ }
+ allMsgs := append(storeMsgs, assistantMsg)
+ if _, err := s.convs.Create(r.Context(), responseID, result.Model, allMsgs); err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to store conversation",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("response_id", responseID),
+ slog.String("error", err.Error()),
+ )...,
+ )
+ // Don't fail the response if storage fails
+ }
+
+ s.logger.InfoContext(r.Context(), "response generated",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("provider", provider.Name()),
+ slog.String("model", result.Model),
+ slog.String("response_id", responseID),
+ slog.Int("input_tokens", result.Usage.InputTokens),
+ slog.Int("output_tokens", result.Usage.OutputTokens),
+ slog.Bool("has_tool_calls", len(result.ToolCalls) > 0),
+ )...,
+ )
+
+ // Build spec-compliant response
+ resp := s.buildResponse(origReq, result, provider.Name(), responseID)
+
+ w.Header().Set("Content-Type", "application/json")
+ w.WriteHeader(http.StatusOK)
+ if err := json.NewEncoder(w).Encode(resp); err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to encode response",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("response_id", responseID),
+ slog.String("error", err.Error()),
+ )...,
+ )
+ }
+}
+
+func (s *GatewayServer) handleStreamingResponse(w http.ResponseWriter, r *http.Request, provider providers.Provider, providerMsgs []api.Message, resolvedReq *api.ResponseRequest, origReq *api.ResponseRequest, storeMsgs []api.Message) {
+ w.Header().Set("Content-Type", "text/event-stream")
+ w.Header().Set("Cache-Control", "no-cache")
+ w.Header().Set("Connection", "keep-alive")
+ w.WriteHeader(http.StatusOK)
+
+ flusher, ok := w.(http.Flusher)
+ if !ok {
+ http.Error(w, "streaming not supported", http.StatusInternalServerError)
+ return
+ }
+
+ responseID := generateID("resp_")
+ itemID := generateID("msg_")
+ seq := 0
+ outputIdx := 0
+ contentIdx := 0
+
+ // Build initial response snapshot (in_progress, no output yet)
+ initialResp := s.buildResponse(origReq, &api.ProviderResult{
+ Model: origReq.Model,
+ }, provider.Name(), responseID)
+ initialResp.Status = "in_progress"
+ initialResp.CompletedAt = nil
+ initialResp.Output = []api.OutputItem{}
+ initialResp.Usage = nil
+
+ // response.created
+ s.sendSSE(w, flusher, &seq, "response.created", &api.StreamEvent{
+ Type: "response.created",
+ Response: initialResp,
+ })
+
+ // response.in_progress
+ s.sendSSE(w, flusher, &seq, "response.in_progress", &api.StreamEvent{
+ Type: "response.in_progress",
+ Response: initialResp,
+ })
+
+ // response.output_item.added
+ inProgressItem := &api.OutputItem{
+ ID: itemID,
+ Type: "message",
+ Status: "in_progress",
+ Role: "assistant",
+ Content: []api.ContentPart{},
+ }
+ s.sendSSE(w, flusher, &seq, "response.output_item.added", &api.StreamEvent{
+ Type: "response.output_item.added",
+ OutputIndex: &outputIdx,
+ Item: inProgressItem,
+ })
+
+ // response.content_part.added
+ emptyPart := &api.ContentPart{
+ Type: "output_text",
+ Text: "",
+ Annotations: []api.Annotation{},
+ }
+ s.sendSSE(w, flusher, &seq, "response.content_part.added", &api.StreamEvent{
+ Type: "response.content_part.added",
+ ItemID: itemID,
+ OutputIndex: &outputIdx,
+ ContentIndex: &contentIdx,
+ Part: emptyPart,
+ })
+
+ // Start provider stream
+ deltaChan, errChan := provider.GenerateStream(r.Context(), providerMsgs, resolvedReq)
+
+ var fullText string
+ var streamErr error
+ var providerModel string
+
+ // Track tool calls being built
+ type toolCallBuilder struct {
+ itemID string
+ id string
+ name string
+ arguments string
+ }
+ toolCallsInProgress := make(map[int]*toolCallBuilder)
+ nextOutputIdx := 0
+ textItemAdded := false
+
+loop:
+ for {
+ select {
+ case delta, ok := <-deltaChan:
+ if !ok {
+ break loop
+ }
+ if delta.Model != "" && providerModel == "" {
+ providerModel = delta.Model
+ }
+
+ // Handle text content
+ if delta.Text != "" {
+ // Add text item on first text delta
+ if !textItemAdded {
+ textItemAdded = true
+ nextOutputIdx++
+ }
+ fullText += delta.Text
+ s.sendSSE(w, flusher, &seq, "response.output_text.delta", &api.StreamEvent{
+ Type: "response.output_text.delta",
+ ItemID: itemID,
+ OutputIndex: &outputIdx,
+ ContentIndex: &contentIdx,
+ Delta: delta.Text,
+ })
+ }
+
+ // Handle tool call delta
+ if delta.ToolCallDelta != nil {
+ tc := delta.ToolCallDelta
+
+ // First chunk for this tool call index
+ if _, exists := toolCallsInProgress[tc.Index]; !exists {
+ toolItemID := generateID("item_")
+ toolOutputIdx := nextOutputIdx
+ nextOutputIdx++
+
+ // Send response.output_item.added
+ s.sendSSE(w, flusher, &seq, "response.output_item.added", &api.StreamEvent{
+ Type: "response.output_item.added",
+ OutputIndex: &toolOutputIdx,
+ Item: &api.OutputItem{
+ ID: toolItemID,
+ Type: "function_call",
+ Status: "in_progress",
+ CallID: tc.ID,
+ Name: tc.Name,
+ },
+ })
+
+ toolCallsInProgress[tc.Index] = &toolCallBuilder{
+ itemID: toolItemID,
+ id: tc.ID,
+ name: tc.Name,
+ arguments: "",
+ }
+ }
+
+ // Send function_call_arguments.delta
+ if tc.Arguments != "" {
+ builder := toolCallsInProgress[tc.Index]
+ builder.arguments += tc.Arguments
+ toolOutputIdx := outputIdx + 1 + tc.Index
+
+ s.sendSSE(w, flusher, &seq, "response.function_call_arguments.delta", &api.StreamEvent{
+ Type: "response.function_call_arguments.delta",
+ ItemID: builder.itemID,
+ OutputIndex: &toolOutputIdx,
+ Delta: tc.Arguments,
+ })
+ }
+ }
+
+ if delta.Done {
+ break loop
+ }
+ case err := <-errChan:
+ if err != nil {
+ streamErr = err
+ }
+ break loop
+ case <-r.Context().Done():
+ s.logger.InfoContext(r.Context(), "client disconnected",
+ slog.String("request_id", logger.FromContext(r.Context())),
+ )
+ return
+ }
+ }
+
+ if streamErr != nil {
+ s.logger.ErrorContext(r.Context(), "stream error",
+ logger.LogAttrsWithTrace(r.Context(),
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("provider", provider.Name()),
+ slog.String("model", origReq.Model),
+ slog.String("error", streamErr.Error()),
+ )...,
+ )
+
+ // Determine error type based on circuit breaker state
+ errorType := "server_error"
+ errorMessage := streamErr.Error()
+ if isCircuitBreakerError(streamErr) {
+ errorType = "circuit_breaker_open"
+ errorMessage = "service temporarily unavailable - circuit breaker open"
+ }
+
+ failedResp := s.buildResponse(origReq, &api.ProviderResult{
+ Model: origReq.Model,
+ }, provider.Name(), responseID)
+ failedResp.Status = "failed"
+ failedResp.CompletedAt = nil
+ failedResp.Output = []api.OutputItem{}
+ failedResp.Error = &api.ResponseError{
+ Type: errorType,
+ Message: errorMessage,
+ }
+ s.sendSSE(w, flusher, &seq, "response.failed", &api.StreamEvent{
+ Type: "response.failed",
+ Response: failedResp,
+ })
+ return
+ }
+
+ // Send done events for text output if text was added
+ if textItemAdded && fullText != "" {
+ // response.output_text.done
+ s.sendSSE(w, flusher, &seq, "response.output_text.done", &api.StreamEvent{
+ Type: "response.output_text.done",
+ ItemID: itemID,
+ OutputIndex: &outputIdx,
+ ContentIndex: &contentIdx,
+ Text: fullText,
+ })
+
+ // response.content_part.done
+ completedPart := &api.ContentPart{
+ Type: "output_text",
+ Text: fullText,
+ Annotations: []api.Annotation{},
+ }
+ s.sendSSE(w, flusher, &seq, "response.content_part.done", &api.StreamEvent{
+ Type: "response.content_part.done",
+ ItemID: itemID,
+ OutputIndex: &outputIdx,
+ ContentIndex: &contentIdx,
+ Part: completedPart,
+ })
+
+ // response.output_item.done
+ completedItem := &api.OutputItem{
+ ID: itemID,
+ Type: "message",
+ Status: "completed",
+ Role: "assistant",
+ Content: []api.ContentPart{*completedPart},
+ }
+ s.sendSSE(w, flusher, &seq, "response.output_item.done", &api.StreamEvent{
+ Type: "response.output_item.done",
+ OutputIndex: &outputIdx,
+ Item: completedItem,
+ })
+ }
+
+ // Send done events for each tool call
+ for idx, builder := range toolCallsInProgress {
+ toolOutputIdx := outputIdx + 1 + idx
+
+ s.sendSSE(w, flusher, &seq, "response.function_call_arguments.done", &api.StreamEvent{
+ Type: "response.function_call_arguments.done",
+ ItemID: builder.itemID,
+ OutputIndex: &toolOutputIdx,
+ Arguments: builder.arguments,
+ })
+
+ s.sendSSE(w, flusher, &seq, "response.output_item.done", &api.StreamEvent{
+ Type: "response.output_item.done",
+ OutputIndex: &toolOutputIdx,
+ Item: &api.OutputItem{
+ ID: builder.itemID,
+ Type: "function_call",
+ Status: "completed",
+ CallID: builder.id,
+ Name: builder.name,
+ Arguments: builder.arguments,
+ },
+ })
+ }
+
+ // Build final completed response
+ model := origReq.Model
+ if providerModel != "" {
+ model = providerModel
+ }
+
+ // Collect tool calls for result
+ var toolCalls []api.ToolCall
+ for _, builder := range toolCallsInProgress {
+ toolCalls = append(toolCalls, api.ToolCall{
+ ID: builder.id,
+ Name: builder.name,
+ Arguments: builder.arguments,
+ })
+ }
+
+ finalResult := &api.ProviderResult{
+ Model: model,
+ Text: fullText,
+ ToolCalls: toolCalls,
+ }
+ completedResp := s.buildResponse(origReq, finalResult, provider.Name(), responseID)
+
+ // Update item IDs to match what we sent during streaming
+ if textItemAdded && len(completedResp.Output) > 0 {
+ completedResp.Output[0].ID = itemID
+ }
+ for idx, builder := range toolCallsInProgress {
+ // Find the corresponding output item
+ for i := range completedResp.Output {
+ if completedResp.Output[i].Type == "function_call" && completedResp.Output[i].CallID == builder.id {
+ completedResp.Output[i].ID = builder.itemID
+ break
+ }
+ }
+ _ = idx // unused
+ }
+
+ // response.completed
+ s.sendSSE(w, flusher, &seq, "response.completed", &api.StreamEvent{
+ Type: "response.completed",
+ Response: completedResp,
+ })
+
+ // Store conversation
+ if fullText != "" || len(toolCalls) > 0 {
+ assistantMsg := api.Message{
+ Role: "assistant",
+ Content: []api.ContentBlock{{Type: "output_text", Text: fullText}},
+ ToolCalls: toolCalls,
+ }
+ allMsgs := append(storeMsgs, assistantMsg)
+ if _, err := s.convs.Create(r.Context(), responseID, model, allMsgs); err != nil {
+ s.logger.ErrorContext(r.Context(), "failed to store conversation",
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("response_id", responseID),
+ slog.String("error", err.Error()),
+ )
+ // Don't fail the response if storage fails
+ }
+
+ s.logger.InfoContext(r.Context(), "streaming response completed",
+ slog.String("request_id", logger.FromContext(r.Context())),
+ slog.String("provider", provider.Name()),
+ slog.String("model", model),
+ slog.String("response_id", responseID),
+ slog.Bool("has_tool_calls", len(toolCalls) > 0),
+ )
+ }
+}
+
+func (s *GatewayServer) sendSSE(w http.ResponseWriter, flusher http.Flusher, seq *int, eventType string, event *api.StreamEvent) {
+ event.SequenceNumber = *seq
+ *seq++
+ data, err := json.Marshal(event)
+ if err != nil {
+ s.logger.Error("failed to marshal SSE event",
+ slog.String("event_type", eventType),
+ slog.String("error", err.Error()),
+ )
+ return
+ }
+ fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, data)
+ flusher.Flush()
+}
+
+func (s *GatewayServer) buildResponse(req *api.ResponseRequest, result *api.ProviderResult, providerName string, responseID string) *api.Response {
+ now := time.Now().Unix()
+
+ model := result.Model
+ if model == "" {
+ model = req.Model
+ }
+
+ // Build output items array
+ outputItems := []api.OutputItem{}
+
+ // Add message item if there's text
+ if result.Text != "" {
+ outputItems = append(outputItems, api.OutputItem{
+ ID: generateID("msg_"),
+ Type: "message",
+ Status: "completed",
+ Role: "assistant",
+ Content: []api.ContentPart{{
+ Type: "output_text",
+ Text: result.Text,
+ Annotations: []api.Annotation{},
+ }},
+ })
+ }
+
+ // Add function_call items
+ for _, tc := range result.ToolCalls {
+ outputItems = append(outputItems, api.OutputItem{
+ ID: generateID("item_"),
+ Type: "function_call",
+ Status: "completed",
+ CallID: tc.ID,
+ Name: tc.Name,
+ Arguments: tc.Arguments,
+ })
+ }
+
+ // Echo back request params with defaults
+ tools := req.Tools
+ if tools == nil {
+ tools = json.RawMessage(`[]`)
+ }
+ toolChoice := req.ToolChoice
+ if toolChoice == nil {
+ toolChoice = json.RawMessage(`"auto"`)
+ }
+ text := req.Text
+ if text == nil {
+ text = json.RawMessage(`{"format":{"type":"text"}}`)
+ }
+ truncation := "disabled"
+ if req.Truncation != nil {
+ truncation = *req.Truncation
+ }
+ temperature := 1.0
+ if req.Temperature != nil {
+ temperature = *req.Temperature
+ }
+ topP := 1.0
+ if req.TopP != nil {
+ topP = *req.TopP
+ }
+ presencePenalty := 0.0
+ if req.PresencePenalty != nil {
+ presencePenalty = *req.PresencePenalty
+ }
+ frequencyPenalty := 0.0
+ if req.FrequencyPenalty != nil {
+ frequencyPenalty = *req.FrequencyPenalty
+ }
+ topLogprobs := 0
+ if req.TopLogprobs != nil {
+ topLogprobs = *req.TopLogprobs
+ }
+ parallelToolCalls := true
+ if req.ParallelToolCalls != nil {
+ parallelToolCalls = *req.ParallelToolCalls
+ }
+ store := true
+ if req.Store != nil {
+ store = *req.Store
+ }
+ background := false
+ if req.Background != nil {
+ background = *req.Background
+ }
+ serviceTier := "default"
+ if req.ServiceTier != nil {
+ serviceTier = *req.ServiceTier
+ }
+ var reasoning json.RawMessage
+ if req.Reasoning != nil {
+ reasoning = req.Reasoning
+ }
+ metadata := req.Metadata
+ if metadata == nil {
+ metadata = map[string]string{}
+ }
+
+ var usage *api.Usage
+ if result.Text != "" {
+ usage = &result.Usage
+ }
+
+ return &api.Response{
+ ID: responseID,
+ Object: "response",
+ CreatedAt: now,
+ CompletedAt: &now,
+ Status: "completed",
+ IncompleteDetails: nil,
+ Model: model,
+ PreviousResponseID: req.PreviousResponseID,
+ Instructions: req.Instructions,
+ Output: outputItems,
+ Error: nil,
+ Tools: tools,
+ ToolChoice: toolChoice,
+ Truncation: truncation,
+ ParallelToolCalls: parallelToolCalls,
+ Text: text,
+ TopP: topP,
+ PresencePenalty: presencePenalty,
+ FrequencyPenalty: frequencyPenalty,
+ TopLogprobs: topLogprobs,
+ Temperature: temperature,
+ Reasoning: reasoning,
+ Usage: usage,
+ MaxOutputTokens: req.MaxOutputTokens,
+ MaxToolCalls: req.MaxToolCalls,
+ Store: store,
+ Background: background,
+ ServiceTier: serviceTier,
+ Metadata: metadata,
+ SafetyIdentifier: nil,
+ PromptCacheKey: nil,
+ Provider: providerName,
+ }
+}
+
+func (s *GatewayServer) resolveProvider(req *api.ResponseRequest) (providers.Provider, error) {
+ if req.Provider != "" {
+ if provider, ok := s.registry.Get(req.Provider); ok {
+ return provider, nil
+ }
+ return nil, fmt.Errorf("provider %s not configured", req.Provider)
+ }
+ return s.registry.Default(req.Model)
+}
+
+func generateID(prefix string) string {
+ id := strings.ReplaceAll(uuid.NewString(), "-", "")
+ return prefix + id[:24]
+}
+
+
+