956 lines
30 KiB
Go
956 lines
30 KiB
Go
package handlers
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
|
|
capspkg "git.gobha.me/xcaliber/chat-switchboard/capabilities"
|
|
"git.gobha.me/xcaliber/chat-switchboard/crypto"
|
|
"git.gobha.me/xcaliber/chat-switchboard/database"
|
|
"git.gobha.me/xcaliber/chat-switchboard/events"
|
|
"git.gobha.me/xcaliber/chat-switchboard/models"
|
|
"git.gobha.me/xcaliber/chat-switchboard/providers"
|
|
"git.gobha.me/xcaliber/chat-switchboard/storage"
|
|
"git.gobha.me/xcaliber/chat-switchboard/store"
|
|
"git.gobha.me/xcaliber/chat-switchboard/tools"
|
|
)
|
|
|
|
// ── Request Types ───────────────────────────
|
|
|
|
type completionRequest struct {
|
|
ChannelID string `json:"channel_id"` // preferred; validated manually below
|
|
ChatID string `json:"chat_id"` // deprecated alias — maps to channel_id
|
|
Content string `json:"content" binding:"required"`
|
|
Model string `json:"model,omitempty"`
|
|
PresetID string `json:"preset_id,omitempty"` // if set, unwraps preset → base model + config
|
|
APIConfigID string `json:"provider_config_id,omitempty"`
|
|
MaxTokens int `json:"max_tokens,omitempty"`
|
|
Temperature *float64 `json:"temperature,omitempty"`
|
|
TopP *float64 `json:"top_p,omitempty"`
|
|
Stream *bool `json:"stream,omitempty"`
|
|
AttachmentIDs []string `json:"attachment_ids,omitempty"` // staged attachment UUIDs to include in request
|
|
DisabledTools []string `json:"disabled_tools,omitempty"` // tool names to exclude from this request
|
|
}
|
|
|
|
// CompletionHandler proxies LLM requests through the backend.
|
|
type CompletionHandler struct {
|
|
vault *crypto.KeyResolver
|
|
stores store.Stores
|
|
hub *events.Hub // WebSocket hub for browser tool bridge
|
|
objStore storage.ObjectStore // file storage for attachment content (nil = disabled)
|
|
}
|
|
|
|
// NewCompletionHandler creates a new handler.
|
|
func NewCompletionHandler(vault *crypto.KeyResolver, stores store.Stores, hub *events.Hub, objStore storage.ObjectStore) *CompletionHandler {
|
|
return &CompletionHandler{vault: vault, stores: stores, hub: hub, objStore: objStore}
|
|
}
|
|
|
|
// ── Chat Completion ─────────────────────────
|
|
// POST /api/v1/chat/completions
|
|
//
|
|
// Flow:
|
|
// 1. Validate request, verify chat ownership
|
|
// 2. Resolve api_config (from request, chat, or user default)
|
|
// 3. Load conversation history from messages
|
|
// 4. Persist user message
|
|
// 5. Attach tool definitions if model supports tools
|
|
// 6. Call provider (stream or non-stream)
|
|
// 7. Tool loop: if LLM requests tool calls, execute them and re-call
|
|
// 8. Stream SSE to client / return JSON
|
|
// 9. Persist assistant message with token counts
|
|
|
|
func (h *CompletionHandler) Complete(c *gin.Context) {
|
|
var req completionRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
// Support chat_id as alias during frontend transition
|
|
channelID := req.ChannelID
|
|
if channelID == "" {
|
|
channelID = req.ChatID
|
|
}
|
|
if channelID == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "channel_id is required"})
|
|
return
|
|
}
|
|
|
|
userID := getUserID(c)
|
|
|
|
// Verify channel ownership
|
|
if !userOwnsChannel(c, channelID, userID) {
|
|
return
|
|
}
|
|
|
|
// ── Preset unwrap: preset overrides defaults, explicit request fields win ──
|
|
var presetSystemPrompt string
|
|
var personaID string // tracks active persona for KB scoping
|
|
if req.PresetID != "" {
|
|
preset := ResolvePreset(h.stores, req.PresetID, userID)
|
|
if preset == nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "preset not found or not accessible"})
|
|
return
|
|
}
|
|
personaID = preset.ID
|
|
// Preset provides defaults; explicit request fields take priority
|
|
if req.Model == "" {
|
|
req.Model = preset.BaseModelID
|
|
}
|
|
if req.APIConfigID == "" && preset.ProviderConfigID != nil {
|
|
req.APIConfigID = *preset.ProviderConfigID
|
|
}
|
|
if req.Temperature == nil && preset.Temperature != nil {
|
|
req.Temperature = preset.Temperature
|
|
}
|
|
if req.MaxTokens == 0 && preset.MaxTokens != nil {
|
|
req.MaxTokens = *preset.MaxTokens
|
|
}
|
|
if preset.SystemPrompt != "" {
|
|
presetSystemPrompt = preset.SystemPrompt
|
|
}
|
|
}
|
|
|
|
// Resolve provider config
|
|
providerCfg, providerID, model, configID, providerScope, err := h.resolveConfig(userID, channelID, req)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
// ── Team policy: require_private_providers ──
|
|
if err := enforcePrivateProviderPolicy(userID, configID); err != nil {
|
|
c.JSON(http.StatusForbidden, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
provider, err := providers.Get(providerID)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
|
|
// Load conversation history
|
|
messages, err := h.loadConversation(channelID, userID, presetSystemPrompt, personaID)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load conversation"})
|
|
return
|
|
}
|
|
|
|
// Resolve capabilities early — needed for vision gating below
|
|
caps := h.getModelCapabilities(c, model, configID)
|
|
|
|
// Build user message — multimodal if attachments are present
|
|
userMsg := providers.Message{
|
|
Role: "user",
|
|
Content: req.Content,
|
|
}
|
|
|
|
if len(req.AttachmentIDs) > 0 && h.objStore != nil {
|
|
parts, augContent, _, err := h.buildMultimodalParts(c, channelID, req.Content, req.AttachmentIDs, caps)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
|
return
|
|
}
|
|
if len(parts) > 0 {
|
|
// Multimodal (images present): use ContentParts
|
|
userMsg.ContentParts = parts
|
|
} else if augContent != req.Content {
|
|
// Doc-only: use augmented text content
|
|
userMsg.Content = augContent
|
|
}
|
|
}
|
|
|
|
messages = append(messages, userMsg)
|
|
|
|
// Persist user message (text-only for storage — multimodal parts are ephemeral)
|
|
msgID, err := h.persistMessage(channelID, userID, "user", req.Content, "", 0, 0, nil, nil)
|
|
if err != nil {
|
|
log.Printf("Failed to persist user message: %v", err)
|
|
}
|
|
|
|
// Link attachments to the persisted message
|
|
if msgID != "" && len(req.AttachmentIDs) > 0 {
|
|
for _, attID := range req.AttachmentIDs {
|
|
if err := h.stores.Attachments.SetMessageID(c.Request.Context(), attID, msgID); err != nil {
|
|
log.Printf("Failed to link attachment %s to message %s: %v", attID, msgID, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Build provider request
|
|
provReq := providers.CompletionRequest{
|
|
Model: model,
|
|
Messages: messages,
|
|
}
|
|
|
|
if req.MaxTokens > 0 {
|
|
provReq.MaxTokens = req.MaxTokens
|
|
} else {
|
|
// ResolveMaxOutput checks: caps → known models → context/8 → 4096
|
|
provReq.MaxTokens = capspkg.ResolveMaxOutput(model, caps)
|
|
}
|
|
|
|
if req.Temperature != nil {
|
|
provReq.Temperature = req.Temperature
|
|
}
|
|
if req.TopP != nil {
|
|
provReq.TopP = req.TopP
|
|
}
|
|
|
|
// Attach tool definitions if model supports tool calling and tools are available
|
|
hasBrowserTools := h.hub != nil && h.hub.IsConnected(userID)
|
|
if caps.ToolCalling && (tools.HasTools() || hasBrowserTools) {
|
|
provReq.Tools = h.buildToolDefs(c.Request.Context(), userID, hasBrowserTools, req.DisabledTools)
|
|
}
|
|
|
|
// Determine streaming
|
|
stream := true // default
|
|
if req.Stream != nil {
|
|
stream = *req.Stream
|
|
}
|
|
|
|
if stream {
|
|
h.streamCompletion(c, provider, providerCfg, provReq, channelID, userID, model, configID, providerScope, personaID)
|
|
} else {
|
|
h.syncCompletion(c, provider, providerCfg, provReq, channelID, userID, model, configID, providerScope, personaID)
|
|
}
|
|
}
|
|
|
|
// buildToolDefs converts registered tools to provider-format tool definitions.
|
|
// If includeBrowser is true, also fetches tool schemas from enabled browser extensions.
|
|
// Tools whose names appear in disabledTools are excluded.
|
|
func (h *CompletionHandler) buildToolDefs(ctx context.Context, userID string, includeBrowser bool, disabledTools []string) []providers.ToolDef {
|
|
// Build disabled set for O(1) lookup
|
|
disabled := make(map[string]bool, len(disabledTools))
|
|
for _, name := range disabledTools {
|
|
disabled[name] = true
|
|
}
|
|
|
|
allTools := tools.AllDefinitionsFiltered(disabled)
|
|
defs := make([]providers.ToolDef, 0, len(allTools))
|
|
for _, t := range allTools {
|
|
defs = append(defs, providers.ToolDef{
|
|
Type: "function",
|
|
Function: providers.FunctionDef{
|
|
Name: t.Name,
|
|
Description: t.Description,
|
|
Parameters: t.Parameters,
|
|
},
|
|
})
|
|
}
|
|
|
|
// Append browser-defined tool schemas from extensions
|
|
if includeBrowser {
|
|
exts, err := h.stores.Extensions.ListForUser(ctx, userID)
|
|
if err != nil {
|
|
log.Printf("⚠️ Failed to load extensions for tools: %v", err)
|
|
return defs
|
|
}
|
|
for _, ext := range exts {
|
|
if ext.Tier != "browser" {
|
|
continue
|
|
}
|
|
var manifest struct {
|
|
Tools []struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
Parameters json.RawMessage `json:"parameters"`
|
|
Tier string `json:"tier"`
|
|
} `json:"tools"`
|
|
}
|
|
if err := json.Unmarshal(ext.Manifest, &manifest); err != nil {
|
|
continue
|
|
}
|
|
for _, t := range manifest.Tools {
|
|
if disabled[t.Name] {
|
|
continue
|
|
}
|
|
defs = append(defs, providers.ToolDef{
|
|
Type: "function",
|
|
Function: providers.FunctionDef{
|
|
Name: t.Name,
|
|
Description: t.Description,
|
|
Parameters: t.Parameters,
|
|
},
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
return defs
|
|
}
|
|
|
|
// ── List Available Tools ────────────────────
|
|
// GET /api/v1/tools
|
|
//
|
|
// Returns all registered server-side tools plus browser extension tools
|
|
// for the current user. Used by the frontend to build the tools toggle menu.
|
|
|
|
type toolInfo struct {
|
|
Name string `json:"name"`
|
|
DisplayName string `json:"display_name"`
|
|
Description string `json:"description"`
|
|
Category string `json:"category"`
|
|
}
|
|
|
|
func (h *CompletionHandler) ListTools(c *gin.Context) {
|
|
userID := c.GetString("user_id")
|
|
|
|
// Server-side tools
|
|
allDefs := tools.AllDefinitions()
|
|
result := make([]toolInfo, 0, len(allDefs))
|
|
for _, d := range allDefs {
|
|
result = append(result, toolInfo{
|
|
Name: d.Name,
|
|
DisplayName: d.DisplayName,
|
|
Description: d.Description,
|
|
Category: d.Category,
|
|
})
|
|
}
|
|
|
|
// Browser extension tools
|
|
if h.hub != nil && h.hub.IsConnected(userID) {
|
|
exts, err := h.stores.Extensions.ListForUser(c.Request.Context(), userID)
|
|
if err == nil {
|
|
for _, ext := range exts {
|
|
if ext.Tier != "browser" {
|
|
continue
|
|
}
|
|
var manifest struct {
|
|
Tools []struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
} `json:"tools"`
|
|
}
|
|
if err := json.Unmarshal(ext.Manifest, &manifest); err != nil {
|
|
continue
|
|
}
|
|
for _, t := range manifest.Tools {
|
|
result = append(result, toolInfo{
|
|
Name: t.Name,
|
|
Description: t.Description,
|
|
Category: "browser",
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
c.JSON(http.StatusOK, gin.H{"data": result})
|
|
}
|
|
|
|
// ── Streaming Completion (SSE) with Tool Loop ──
|
|
|
|
const maxToolIterations = 10 // safety limit on tool call rounds
|
|
|
|
func (h *CompletionHandler) streamCompletion(
|
|
c *gin.Context,
|
|
provider providers.Provider,
|
|
cfg providers.ProviderConfig,
|
|
req providers.CompletionRequest,
|
|
channelID, userID, model, configID, providerScope, personaID string,
|
|
) {
|
|
result := streamWithToolLoop(c, provider, cfg, &req, model, userID, channelID, personaID, h.hub)
|
|
|
|
// Persist assistant response
|
|
if result.Content != "" {
|
|
if _, err := h.persistMessage(channelID, userID, "assistant", result.Content, model, result.InputTokens, result.OutputTokens, nil, toolActivityJSON(result.ToolActivity)); err != nil {
|
|
log.Printf("Failed to persist assistant message: %v", err)
|
|
}
|
|
}
|
|
|
|
// Log usage
|
|
h.logUsage(c, channelID, userID, configID, providerScope, model,
|
|
result.InputTokens, result.OutputTokens,
|
|
result.CacheCreationTokens, result.CacheReadTokens)
|
|
}
|
|
|
|
// ── Non-Streaming Completion with Tool Loop ──
|
|
|
|
func (h *CompletionHandler) syncCompletion(
|
|
c *gin.Context,
|
|
provider providers.Provider,
|
|
cfg providers.ProviderConfig,
|
|
req providers.CompletionRequest,
|
|
channelID, userID, model, configID, providerScope, personaID string,
|
|
) {
|
|
var totalInput, totalOutput int
|
|
var totalCacheCreation, totalCacheRead int
|
|
var finalContent string
|
|
var allToolActivity []map[string]interface{}
|
|
|
|
for iteration := 0; iteration < maxToolIterations; iteration++ {
|
|
resp, err := provider.ChatCompletion(c.Request.Context(), cfg, req)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadGateway, gin.H{"error": "provider error: " + err.Error()})
|
|
return
|
|
}
|
|
|
|
totalInput += resp.InputTokens
|
|
totalOutput += resp.OutputTokens
|
|
totalCacheCreation += resp.CacheCreationTokens
|
|
totalCacheRead += resp.CacheReadTokens
|
|
|
|
// No tool calls — normal response
|
|
if len(resp.ToolCalls) == 0 || resp.FinishReason != "tool_calls" {
|
|
finalContent = resp.Content
|
|
|
|
// Persist assistant response
|
|
if _, err := h.persistMessage(channelID, userID, "assistant", finalContent, model, totalInput, totalOutput, nil, toolActivityJSON(allToolActivity)); err != nil {
|
|
log.Printf("Failed to persist assistant message: %v", err)
|
|
}
|
|
|
|
// Log usage
|
|
h.logUsage(c, channelID, userID, configID, providerScope, model,
|
|
totalInput, totalOutput, totalCacheCreation, totalCacheRead)
|
|
|
|
// Return OpenAI-compatible response
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"model": resp.Model,
|
|
"choices": []gin.H{
|
|
{
|
|
"message": gin.H{
|
|
"role": "assistant",
|
|
"content": finalContent,
|
|
},
|
|
"finish_reason": resp.FinishReason,
|
|
},
|
|
},
|
|
"usage": gin.H{
|
|
"prompt_tokens": totalInput,
|
|
"completion_tokens": totalOutput,
|
|
"total_tokens": totalInput + totalOutput,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// ── Tool execution round ────────────────
|
|
assistantMsg := providers.Message{
|
|
Role: "assistant",
|
|
Content: resp.Content,
|
|
}
|
|
for _, tc := range resp.ToolCalls {
|
|
assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, tc)
|
|
}
|
|
req.Messages = append(req.Messages, assistantMsg)
|
|
|
|
execCtx := tools.ExecutionContext{
|
|
UserID: userID,
|
|
ChannelID: channelID,
|
|
PersonaID: personaID,
|
|
}
|
|
for _, tc := range resp.ToolCalls {
|
|
call := tools.ToolCall{
|
|
ID: tc.ID,
|
|
Name: tc.Function.Name,
|
|
Arguments: tc.Function.Arguments,
|
|
}
|
|
log.Printf("🔧 Executing tool: %s (call %s)", call.Name, call.ID)
|
|
result := tools.ExecuteCall(c.Request.Context(), execCtx, call)
|
|
|
|
// Collect for persistence
|
|
allToolActivity = append(allToolActivity, map[string]interface{}{
|
|
"id": tc.ID,
|
|
"name": tc.Function.Name,
|
|
"arguments": tc.Function.Arguments,
|
|
"result": result.Content,
|
|
"is_error": result.IsError,
|
|
})
|
|
|
|
req.Messages = append(req.Messages, providers.Message{
|
|
Role: "tool",
|
|
Content: result.Content,
|
|
ToolCallID: result.ToolCallID,
|
|
Name: result.Name,
|
|
})
|
|
}
|
|
}
|
|
|
|
// Max iterations reached
|
|
c.JSON(http.StatusOK, gin.H{
|
|
"model": model,
|
|
"choices": []gin.H{
|
|
{
|
|
"message": gin.H{
|
|
"role": "assistant",
|
|
"content": "[Tool execution limit reached]",
|
|
},
|
|
"finish_reason": "stop",
|
|
},
|
|
},
|
|
})
|
|
}
|
|
|
|
// ── Model Capabilities ──────────────────────
|
|
|
|
// escapeJSON escapes a string for safe embedding in a JSON string value.
|
|
func escapeJSON(s string) string {
|
|
s = strings.ReplaceAll(s, `\`, `\\`)
|
|
s = strings.ReplaceAll(s, `"`, `\"`)
|
|
s = strings.ReplaceAll(s, "\n", `\n`)
|
|
s = strings.ReplaceAll(s, "\r", `\r`)
|
|
s = strings.ReplaceAll(s, "\t", `\t`)
|
|
return s
|
|
}
|
|
|
|
// getModelCapabilities looks up capabilities from model_catalog DB,
|
|
// then overlays with known model defaults and heuristic detection.
|
|
func (h *CompletionHandler) getModelCapabilities(c *gin.Context, model, apiConfigID string) models.ModelCapabilities {
|
|
return ResolveModelCaps(c, model, apiConfigID)
|
|
}
|
|
|
|
// ── Multimodal Assembly ─────────────────────
|
|
// Builds content parts from attachments for the user message.
|
|
//
|
|
// Returns:
|
|
// - parts: ContentParts array (non-nil only when images are present)
|
|
// - augContent: enriched text content with document context (for doc-only case)
|
|
// - validIDs: attachment IDs that were successfully processed
|
|
// - error: if any attachment is invalid or vision is needed but missing
|
|
//
|
|
// Rules:
|
|
// - Images → base64 data URI (requires vision capability)
|
|
// - Documents with extracted_text → text injection
|
|
// - Documents without extraction → filename placeholder
|
|
// - Text is always the first part
|
|
|
|
func (h *CompletionHandler) buildMultimodalParts(
|
|
c *gin.Context,
|
|
channelID, textContent string,
|
|
attachmentIDs []string,
|
|
caps models.ModelCapabilities,
|
|
) ([]providers.ContentPart, string, []string, error) {
|
|
|
|
parts := []providers.ContentPart{
|
|
{Type: "text", Text: textContent},
|
|
}
|
|
var docTexts []string
|
|
var validIDs []string
|
|
hasImage := false
|
|
|
|
for _, attID := range attachmentIDs {
|
|
att, err := h.stores.Attachments.GetByID(c.Request.Context(), attID)
|
|
if err != nil {
|
|
return nil, "", nil, fmt.Errorf("attachment %s not found", attID)
|
|
}
|
|
|
|
// Security: verify attachment belongs to this channel
|
|
if att.ChannelID != channelID {
|
|
return nil, "", nil, fmt.Errorf("attachment %s does not belong to this channel", attID)
|
|
}
|
|
|
|
if isImageContentType(att.ContentType) {
|
|
// Vision gating: reject images if model lacks vision
|
|
if !caps.Vision {
|
|
return nil, "", nil, fmt.Errorf("model does not support image input; remove image attachments or choose a vision-capable model")
|
|
}
|
|
|
|
// Read image from storage, base64 encode, build data URI
|
|
reader, _, _, err := h.objStore.Get(c.Request.Context(), att.StorageKey)
|
|
if err != nil {
|
|
log.Printf("Failed to read attachment %s from storage: %v", attID, err)
|
|
parts = append(parts, providers.ContentPart{
|
|
Type: "text",
|
|
Text: fmt.Sprintf("[Image: %s — failed to read from storage]", att.Filename),
|
|
})
|
|
validIDs = append(validIDs, attID)
|
|
continue
|
|
}
|
|
data, err := io.ReadAll(reader)
|
|
reader.Close()
|
|
if err != nil {
|
|
log.Printf("Failed to read attachment %s bytes: %v", attID, err)
|
|
validIDs = append(validIDs, attID)
|
|
continue
|
|
}
|
|
|
|
b64 := base64.StdEncoding.EncodeToString(data)
|
|
dataURI := fmt.Sprintf("data:%s;base64,%s", att.ContentType, b64)
|
|
|
|
parts = append(parts, providers.ContentPart{
|
|
Type: "image_url",
|
|
ImageURL: &providers.ImageURL{
|
|
URL: dataURI,
|
|
Detail: "auto",
|
|
},
|
|
})
|
|
hasImage = true
|
|
|
|
} else if att.ExtractedText != nil && *att.ExtractedText != "" {
|
|
// Document with extracted text → inject as context
|
|
docText := fmt.Sprintf("[Document: %s]\n%s", att.Filename, *att.ExtractedText)
|
|
parts = append(parts, providers.ContentPart{
|
|
Type: "text",
|
|
Text: docText,
|
|
})
|
|
docTexts = append(docTexts, docText)
|
|
|
|
} else {
|
|
// Document without extraction (pending, failed, or not extractable)
|
|
status := "pending"
|
|
if s, ok := att.Metadata["extraction_status"].(string); ok {
|
|
status = s
|
|
}
|
|
placeholder := fmt.Sprintf("[Attached file: %s (extraction %s)]", att.Filename, status)
|
|
parts = append(parts, providers.ContentPart{
|
|
Type: "text",
|
|
Text: placeholder,
|
|
})
|
|
docTexts = append(docTexts, placeholder)
|
|
}
|
|
|
|
validIDs = append(validIDs, attID)
|
|
}
|
|
|
|
// If images present → use ContentParts (multimodal array)
|
|
if hasImage {
|
|
return parts, textContent, validIDs, nil
|
|
}
|
|
|
|
// Doc-only: merge document context into a single text string (more efficient)
|
|
augmented := textContent
|
|
for _, dt := range docTexts {
|
|
augmented += "\n\n" + dt
|
|
}
|
|
return nil, augmented, validIDs, nil
|
|
}
|
|
|
|
// isImageContentType returns true for MIME types that should be sent as
|
|
// base64 image content parts rather than text extraction.
|
|
func isImageContentType(ct string) bool {
|
|
return strings.HasPrefix(ct, "image/")
|
|
}
|
|
|
|
// ── Config Resolution ───────────────────────
|
|
// Priority: request.provider_config_id → chat.provider_config_id → user's first active config
|
|
|
|
func (h *CompletionHandler) resolveConfig(userID string, channelID string, req completionRequest) (providers.ProviderConfig, string, string, string, string, error) {
|
|
var configID string
|
|
|
|
// 1. Explicit config from request
|
|
if req.APIConfigID != "" {
|
|
configID = req.APIConfigID
|
|
}
|
|
|
|
// 2. Config from channel
|
|
if configID == "" {
|
|
var channelConfigID *string
|
|
err := database.DB.QueryRow(
|
|
`SELECT provider_config_id FROM channels WHERE id = $1`, channelID,
|
|
).Scan(&channelConfigID)
|
|
if err == nil && channelConfigID != nil {
|
|
configID = *channelConfigID
|
|
}
|
|
}
|
|
|
|
// 3. User's first active config (personal first, then global — excludes team providers)
|
|
if configID == "" {
|
|
err := database.DB.QueryRow(`
|
|
SELECT id FROM provider_configs
|
|
WHERE is_active = true AND (
|
|
(scope = 'personal' AND owner_id = $1)
|
|
OR scope = 'global'
|
|
)
|
|
ORDER BY scope ASC, created_at ASC
|
|
LIMIT 1
|
|
`, userID).Scan(&configID)
|
|
if err != nil {
|
|
return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("no API config found — add one at /api-configs")
|
|
}
|
|
}
|
|
|
|
// Load the config — allow personal, global, OR team configs the user belongs to
|
|
var providerID, endpoint string
|
|
var providerScope string
|
|
var modelDefault *string
|
|
var apiKeyEnc, keyNonce []byte
|
|
var keyScope string
|
|
var customHeadersJSON, providerSettingsJSON []byte
|
|
err := database.DB.QueryRow(`
|
|
SELECT provider, endpoint, scope, api_key_enc, key_nonce, key_scope,
|
|
model_default, headers, settings
|
|
FROM provider_configs
|
|
WHERE id = $1 AND is_active = true
|
|
AND (scope = 'global'
|
|
OR (scope = 'personal' AND owner_id = $2)
|
|
OR (scope = 'team' AND owner_id IN (SELECT team_id FROM team_members WHERE user_id = $2)))
|
|
`, configID, userID).Scan(&providerID, &endpoint, &providerScope, &apiKeyEnc, &keyNonce, &keyScope,
|
|
&modelDefault, &customHeadersJSON, &providerSettingsJSON)
|
|
|
|
if err == sql.ErrNoRows {
|
|
return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("API config not found or not accessible")
|
|
}
|
|
if err != nil {
|
|
return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("failed to load API config: %w", err)
|
|
}
|
|
|
|
// Resolve model: request > config default
|
|
model := req.Model
|
|
if model == "" && modelDefault != nil {
|
|
model = *modelDefault
|
|
}
|
|
if model == "" {
|
|
return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("no model specified and no default model in config")
|
|
}
|
|
|
|
// Decrypt API key using the appropriate tier
|
|
key := ""
|
|
if len(apiKeyEnc) > 0 {
|
|
if h.vault != nil {
|
|
var err error
|
|
key, err = h.vault.Decrypt(apiKeyEnc, keyNonce, keyScope, userID)
|
|
if err != nil {
|
|
if err == crypto.ErrVaultLocked {
|
|
return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("personal vault is locked — please log in again")
|
|
}
|
|
return providers.ProviderConfig{}, "", "", "", "", fmt.Errorf("failed to decrypt API key: %w", err)
|
|
}
|
|
} else {
|
|
// No vault — key stored as raw bytes (unencrypted fallback)
|
|
key = string(apiKeyEnc)
|
|
}
|
|
}
|
|
|
|
// Parse custom headers
|
|
customHeaders := make(map[string]string)
|
|
if customHeadersJSON != nil {
|
|
_ = json.Unmarshal(customHeadersJSON, &customHeaders)
|
|
}
|
|
|
|
// Parse provider-specific settings
|
|
providerSettings := make(map[string]interface{})
|
|
if providerSettingsJSON != nil {
|
|
_ = json.Unmarshal(providerSettingsJSON, &providerSettings)
|
|
}
|
|
|
|
return providers.ProviderConfig{
|
|
Endpoint: endpoint,
|
|
APIKey: key,
|
|
CustomHeaders: customHeaders,
|
|
Settings: providerSettings,
|
|
}, providerID, model, configID, providerScope, nil
|
|
}
|
|
|
|
// ── Conversation Loader ─────────────────────
|
|
|
|
// loadConversation builds the message history for the LLM by walking the
|
|
// active path through the message tree. Only the current branch is sent —
|
|
// sibling branches are never included in context.
|
|
//
|
|
// Summary-aware: if the path contains a summary node (metadata.type = "summary"),
|
|
// messages before it are replaced by the summary content as a system message.
|
|
func (h *CompletionHandler) loadConversation(channelID, userID, presetSystemPrompt, personaID string) ([]providers.Message, error) {
|
|
messages := make([]providers.Message, 0)
|
|
|
|
// ── Admin system prompt (always injected first, no opt out) ──
|
|
adminPrompt, err := h.stores.GlobalConfig.Get(context.Background(), "system_prompt")
|
|
if err == nil {
|
|
if content, ok := adminPrompt["content"].(string); ok && content != "" {
|
|
messages = append(messages, providers.Message{
|
|
Role: "system",
|
|
Content: content,
|
|
})
|
|
}
|
|
}
|
|
|
|
// ── User/preset system prompt (appended after admin prompt) ──
|
|
var systemPrompt *string
|
|
_ = database.DB.QueryRow(
|
|
`SELECT system_prompt FROM channels WHERE id = $1`, channelID,
|
|
).Scan(&systemPrompt)
|
|
|
|
// Preset system prompt takes priority; channel system prompt is fallback
|
|
if presetSystemPrompt != "" {
|
|
messages = append(messages, providers.Message{
|
|
Role: "system",
|
|
Content: presetSystemPrompt,
|
|
})
|
|
} else if systemPrompt != nil && *systemPrompt != "" {
|
|
messages = append(messages, providers.Message{
|
|
Role: "system",
|
|
Content: *systemPrompt,
|
|
})
|
|
}
|
|
|
|
// ── KB hint (nudge LLM to use kb_search when KBs are active) ──
|
|
if kbHint := BuildKBHint(context.Background(), h.stores, channelID, userID, personaID); kbHint != "" {
|
|
messages = append(messages, providers.Message{
|
|
Role: "system",
|
|
Content: kbHint,
|
|
})
|
|
}
|
|
|
|
// Walk the active path (root → leaf) instead of loading all messages
|
|
path, err := getActivePath(channelID, userID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// Find the most recent summary node in the path
|
|
summaryIdx := -1
|
|
for i, m := range path {
|
|
if isSummaryMessage(&m) {
|
|
summaryIdx = i
|
|
}
|
|
}
|
|
|
|
// If we found a summary, inject it as a system message and skip prior messages
|
|
startIdx := 0
|
|
if summaryIdx >= 0 {
|
|
messages = append(messages, providers.Message{
|
|
Role: "system",
|
|
Content: "Previous conversation summary:\n" + path[summaryIdx].Content,
|
|
})
|
|
startIdx = summaryIdx + 1
|
|
}
|
|
|
|
for _, m := range path[startIdx:] {
|
|
if m.Role == "system" {
|
|
continue // system prompts handled above
|
|
}
|
|
// Skip summary nodes that aren't the boundary (shouldn't happen, but defensive)
|
|
if isSummaryMessage(&m) {
|
|
continue
|
|
}
|
|
messages = append(messages, providers.Message{
|
|
Role: m.Role,
|
|
Content: m.Content,
|
|
})
|
|
}
|
|
|
|
return messages, nil
|
|
}
|
|
|
|
// ── Message Persistence ─────────────────────
|
|
|
|
// persistMessage inserts a message into the tree, updates the cursor, and
|
|
// returns the new message's ID.
|
|
//
|
|
// Parent resolution:
|
|
// - parentOverride (explicit parent for edit/regen operations)
|
|
// - cursor's active_leaf_id (normal conversation flow)
|
|
// - latest message by created_at (fallback for legacy/missing cursor)
|
|
// toolActivityJSON marshals tool activity for persistence. Returns nil if empty.
|
|
func toolActivityJSON(activity []map[string]interface{}) json.RawMessage {
|
|
if len(activity) == 0 {
|
|
return nil
|
|
}
|
|
b, _ := json.Marshal(activity)
|
|
return b
|
|
}
|
|
|
|
func (h *CompletionHandler) persistMessage(channelID, userID, role, content, model string, inputTokens, outputTokens int, parentOverride *string, toolCallsJSON json.RawMessage) (string, error) {
|
|
var tokensUsed *int
|
|
if inputTokens > 0 || outputTokens > 0 {
|
|
total := inputTokens + outputTokens
|
|
tokensUsed = &total
|
|
}
|
|
|
|
// Determine parent
|
|
var parentID *string
|
|
if parentOverride != nil {
|
|
parentID = parentOverride
|
|
} else {
|
|
parentID, _ = getActiveLeaf(channelID, userID)
|
|
}
|
|
|
|
// Compute sibling_index
|
|
siblingIdx := nextSiblingIndex(channelID, parentID)
|
|
|
|
// Determine participant
|
|
participantType := "user"
|
|
participantID := userID
|
|
if role == "assistant" {
|
|
participantType = "model"
|
|
participantID = model
|
|
}
|
|
|
|
// Prepare tool_calls JSONB (nil → NULL)
|
|
var tcVal interface{}
|
|
if len(toolCallsJSON) > 0 {
|
|
tcVal = string(toolCallsJSON)
|
|
}
|
|
|
|
// Insert with RETURNING id
|
|
var newID string
|
|
err := database.DB.QueryRow(`
|
|
INSERT INTO messages (channel_id, role, content, model, tokens_used,
|
|
tool_calls, parent_id, participant_type, participant_id, sibling_index)
|
|
VALUES ($1, $2, $3, NULLIF($4, ''), $5, $6, $7, $8, $9, $10)
|
|
RETURNING id
|
|
`, channelID, role, content, model, tokensUsed,
|
|
tcVal, parentID, participantType, participantID, siblingIdx,
|
|
).Scan(&newID)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// Update cursor to point at new message
|
|
if userID != "" {
|
|
_ = updateCursor(channelID, userID, newID)
|
|
}
|
|
|
|
// Touch channel updated_at
|
|
_, _ = database.DB.Exec(`UPDATE channels SET updated_at = NOW() WHERE id = $1`, channelID)
|
|
return newID, nil
|
|
}
|
|
|
|
// ── Usage Logging ─────────────────────────
|
|
|
|
// logUsage records token usage and cost in the usage_log table.
|
|
// Always logs even if tokens are zero — the request happened.
|
|
func (h *CompletionHandler) logUsage(
|
|
c *gin.Context,
|
|
channelID, userID, configID, providerScope, modelID string,
|
|
inputTokens, outputTokens, cacheCreation, cacheRead int,
|
|
) {
|
|
if h.stores.Usage == nil {
|
|
return
|
|
}
|
|
|
|
entry := &models.UsageEntry{
|
|
ChannelID: &channelID,
|
|
UserID: userID,
|
|
ProviderConfigID: &configID,
|
|
ProviderScope: providerScope,
|
|
ModelID: modelID,
|
|
PromptTokens: inputTokens,
|
|
CompletionTokens: outputTokens,
|
|
CacheCreationTokens: cacheCreation,
|
|
CacheReadTokens: cacheRead,
|
|
}
|
|
|
|
// Look up pricing for cost calculation
|
|
if h.stores.Pricing != nil {
|
|
pricing, err := h.stores.Pricing.GetForModel(c.Request.Context(), configID, modelID)
|
|
if err == nil && pricing != nil {
|
|
entry.CostInput = calcCost(inputTokens, pricing.InputPerM)
|
|
entry.CostOutput = calcCost(outputTokens, pricing.OutputPerM)
|
|
}
|
|
}
|
|
|
|
if err := h.stores.Usage.Log(c.Request.Context(), entry); err != nil {
|
|
log.Printf("⚠ Failed to log usage: %v", err)
|
|
}
|
|
}
|
|
|
|
// calcCost computes the cost for a given token count and price-per-million.
|
|
func calcCost(tokens int, pricePerM *float64) *float64 {
|
|
if pricePerM == nil || tokens == 0 {
|
|
return nil
|
|
}
|
|
cost := float64(tokens) * *pricePerM / 1_000_000.0
|
|
return &cost
|
|
}
|