llama-link/internal/api/handlers/chat.go

package handlers

import (
	"encoding/json"
	"log/slog"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"

	"github.com/llamalink/llamalink/internal/api/middleware"
	"github.com/llamalink/llamalink/internal/auth"
	"github.com/llamalink/llamalink/internal/db"
	"github.com/llamalink/llamalink/internal/llama"
	"github.com/llamalink/llamalink/internal/quota"
)

type ChatHandler struct {
	proxy       *llama.Proxy
	authService *auth.Service
	quotaSvc    *quota.Service
	webhookSvc  *quota.WebhookService
}

func NewChatHandler(proxy *llama.Proxy, authService *auth.Service, quotaSvc *quota.Service, webhookSvc *quota.WebhookService) *ChatHandler {
	return &ChatHandler{
		proxy:       proxy,
		authService: authService,
		quotaSvc:    quotaSvc,
		webhookSvc:  webhookSvc,
	}
}

type ChatCompletionRequest struct {
	Model    string                 `json:"model" binding:"required"`
	Messages []llama.ChatMessage    `json:"messages" binding:"required"`
	Stream   bool                   `json:"stream"`
	MaxTokens int                   `json:"max_tokens"`
	Temperature float64             `json:"temperature"`
	TopP      float64               `json:"top_p"`
}

func (h *ChatHandler) ChatCompletions(c *gin.Context) {
	start := time.Now()

	var req ChatCompletionRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{
			"error": gin.H{
				"code":    "validation_error",
				"message": err.Error(),
			},
		})
		return
	}

	apiKey := middleware.GetAPIKey(c)

	// Check quota
	ok, msg, err := h.quotaSvc.CheckQuota(apiKey.ID, req.Model)
	if err != nil {
		slog.Error("quota check failed", "error", err)
	}
	if !ok {
		h.logUsage(c, apiKey, req.Model, 0, 0, 0, "quota_exceeded", start)
		h.webhookSvc.Dispatch("quota_exceeded", apiKey, msg, nil)
		c.JSON(http.StatusTooManyRequests, gin.H{
			"error": gin.H{
				"code":             "quota_exceeded",
				"message":          msg,
				"retry_after_seconds": nil,
			},
		})
		return
	}

	// Check model readiness
	if !h.proxy.Manager().IsReady() {
		c.JSON(http.StatusServiceUnavailable, gin.H{
			"error": gin.H{
				"code":    "model_not_loaded",
				"message": "Model not ready, use POST /v1/models/{name}/load",
			},
		})
		return
	}

	// Check rate limit (basic)
	// TODO: implement token bucket

	if req.Stream {
		h.handleStream(c, apiKey, req, start)
		return
	}

	// Non-streaming
	resp, err := h.proxy.ChatCompletion(c.Request.Context(), llama.ChatCompletionRequest{
		Model:       req.Model,
		Messages:    req.Messages,
		MaxTokens:   req.MaxTokens,
		Temperature: req.Temperature,
		TopP:        req.TopP,
	})

	if err != nil {
		h.logUsage(c, apiKey, req.Model, 0, 0, 0, "error", start)
		c.JSON(http.StatusBadGateway, gin.H{
			"error": gin.H{
				"code":    "upstream_error",
				"message": err.Error(),
			},
		})
		return
	}

	// Consume quota
	totalTokens := resp.Usage.TotalTokens
	if err := h.quotaSvc.ConsumeQuota(apiKey.ID, req.Model, totalTokens); err != nil && err != quota.ErrQuotaExceeded {
		slog.Error("failed to consume quota", "error", err)
	}

	h.logUsage(c, apiKey, req.Model, resp.Usage.PromptTokens, resp.Usage.CompletionTokens, totalTokens, "success", start)

	// Convert to OpenAI format
	c.JSON(http.StatusOK, gin.H{
		"id":      resp.ID,
		"object":  "chat.completion",
		"created": resp.Created,
		"model":   resp.Model,
		"choices": []gin.H{{
			"index": 0,
			"message": gin.H{
				"role":    resp.Choices[0].Message.Role,
				"content": resp.Choices[0].Message.Content,
			},
			"finish_reason": resp.Choices[0].FinishReason,
		}},
		"usage": gin.H{
			"prompt_tokens":     resp.Usage.PromptTokens,
			"completion_tokens": resp.Usage.CompletionTokens,
			"total_tokens":      resp.Usage.TotalTokens,
		},
	})
}

func (h *ChatHandler) handleStream(c *gin.Context, apiKey *db.ApiKey, req ChatCompletionRequest, start time.Time) {
	stream, errCh := h.proxy.ChatCompletionStream(c.Request.Context(), llama.ChatCompletionRequest{
		Model:       req.Model,
		Messages:    req.Messages,
		MaxTokens:   req.MaxTokens,
		Temperature: req.Temperature,
		TopP:        req.TopP,
	})

	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("Transfer-Encoding", "chunked")

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "streaming not supported"})
		return
	}

	totalTokens := 0
	promptTokens := 0
	completionTokens := 0

	for {
		select {
		case resp, ok := <-stream:
			if !ok {
				flusher.Flush()
				return
			}

			// Count tokens roughly
			completionTokens += len(resp.Choices[0].Delta.Content) / 4

			// Write SSE
			_, _ = c.Writer.WriteString("data: ")
			_, _ = c.Writer.WriteString("{\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"created\":")
			_, _ = c.Writer.WriteString(formatInt(resp.Created))
			_, _ = c.Writer.WriteString(",\"model\":\"")
			_, _ = c.Writer.WriteString(resp.Model)
			_, _ = c.Writer.WriteString("\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"")
			_, _ = c.Writer.WriteString(escapeJSON(resp.Choices[0].Delta.Content))
			_, _ = c.Writer.WriteString("\"}}]}\n\n")
			flusher.Flush()

		case err := <-errCh:
			h.logUsage(c, apiKey, req.Model, promptTokens, completionTokens, totalTokens, "error", start)
			_, _ = c.Writer.WriteString("data: [DONE]\n\n")
			flusher.Flush()
			if err != nil {
				slog.Error("stream error", "error", err)
			}
			return

		case <-c.Request.Context().Done():
			return
		}
	}
}

func (h *ChatHandler) logUsage(c *gin.Context, apiKey *db.ApiKey, model string, promptTokens, completionTokens, totalTokens int, status string, start time.Time) {
	latencyMs := int(time.Since(start).Milliseconds())

	log := &db.UsageLog{
		ID:              uuid.New(),
		ApiKeyID:        apiKey.ID,
		ModelName:       model,
		Endpoint:        "/v1/chat/completions",
		PromptTokens:    promptTokens,
		CompletionTokens: completionTokens,
		TotalTokens:     totalTokens,
		LatencyMs:       latencyMs,
		Status:          status,
		Streamed:        false,
	}

	if ip := c.ClientIP(); ip != "" {
		log.IPAddress = &ip
	}
	if ua := c.GetHeader("User-Agent"); ua != "" {
		log.UserAgent = &ua
	}

	// Async log
	go func() {
		// Would use a separate goroutine-safe session here
	}()
	_ = log // avoid unused warning
}

func (h *ChatHandler) Manager() *llama.Manager {
	return h.proxy.Manager()
}

func formatInt(n int64) string {
	return string(rune(n))
}

func escapeJSON(s string) string {
	b, _ := json.Marshal(s)
	return string(b[1 : len(b)-1])
}