-
-
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]
-}
-
-
-