core/server/tools/kbsearch.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"log"

	"git.gobha.me/xcaliber/chat-switchboard/knowledge"
	"git.gobha.me/xcaliber/chat-switchboard/store"
)

// ── Late Registration ────────────────────────
// kb_search cannot use init() because it needs stores + embedder,
// which aren't available until after main.go initializes them.
// Call RegisterKBSearch() from main.go after stores init.

// RegisterKBSearch registers the kb_search tool with captured dependencies.
func RegisterKBSearch(stores store.Stores, embedder *knowledge.Embedder) {
	Register(&kbSearchTool{
		stores:   stores,
		embedder: embedder,
	})
}

// ═══════════════════════════════════════════
// kb_search
// ═══════════════════════════════════════════

type kbSearchTool struct {
	stores   store.Stores
	embedder *knowledge.Embedder
}

func (t *kbSearchTool) Definition() ToolDef {
	return ToolDef{
		Name:        "kb_search",
		DisplayName: "Knowledge Base",
		Category:    "knowledge",
		Description: "Search knowledge bases for relevant information. " +
			"Returns text passages from uploaded documents that match the query. " +
			"Use this when the user asks questions that might be answered by their documents.",
		Parameters: JSONSchema(map[string]interface{}{
			"query": Prop("string", "Search query — use natural language to describe what you're looking for"),
			"max_results": map[string]interface{}{
				"type":        "integer",
				"description": "Maximum results to return (1-20, default 5)",
			},
		}, []string{"query"}),
	}
}

func (t *kbSearchTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
	var args struct {
		Query      string `json:"query"`
		MaxResults int    `json:"max_results"`
	}
	if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
		return "", fmt.Errorf("invalid arguments: %w", err)
	}

	if args.Query == "" {
		return "", fmt.Errorf("query is required")
	}
	if args.MaxResults <= 0 || args.MaxResults > 20 {
		args.MaxResults = 5
	}

	// Resolve user's team IDs for scoped access
	teamIDs, _ := t.stores.Teams.GetUserTeamIDs(ctx, execCtx.UserID)

	// Get active KB IDs for this channel, including persona-bound KBs
	var kbIDs []string
	var err error
	if execCtx.PersonaID != "" {
		kbIDs, err = t.stores.KnowledgeBases.GetActiveKBIDsWithPersona(
			ctx, execCtx.ChannelID, execCtx.UserID, teamIDs, execCtx.PersonaID)
	} else {
		kbIDs, err = t.stores.KnowledgeBases.GetActiveKBIDs(
			ctx, execCtx.ChannelID, execCtx.UserID, teamIDs)
	}
	if err != nil {
		return "", fmt.Errorf("failed to look up active knowledge bases: %w", err)
	}

	// Also include project-bound KBs (v0.19.0)
	if t.stores.Projects != nil {
		projID, _ := t.stores.Projects.GetProjectIDForChannel(ctx, execCtx.ChannelID)
		if projID != "" {
			projKBIDs, _ := t.stores.Projects.GetKBIDs(ctx, projID)
			projSet := make(map[string]bool, len(kbIDs))
			for _, id := range kbIDs {
				projSet[id] = true
			}
			for _, id := range projKBIDs {
				if !projSet[id] {
					kbIDs = append(kbIDs, id)
					projSet[id] = true
				}
			}
		}
	}

	// Also include personal KBs (always available to owner, even if not linked to channel)
	personalKBs, err := t.stores.KnowledgeBases.ListPersonal(ctx, execCtx.UserID)
	if err == nil {
		kbIDSet := make(map[string]bool, len(kbIDs))
		for _, id := range kbIDs {
			kbIDSet[id] = true
		}
		for _, kb := range personalKBs {
			if !kbIDSet[kb.ID] && kb.ChunkCount > 0 {
				kbIDs = append(kbIDs, kb.ID)
			}
		}
	}

	if len(kbIDs) == 0 {
		result, _ := json.Marshal(map[string]interface{}{
			"results":      []interface{}{},
			"query":        args.Query,
			"searched_kbs": []string{},
			"message":      "No knowledge bases are active on this channel. Link a knowledge base to the channel or upload documents to a personal knowledge base first.",
		})
		return string(result), nil
	}

	// Resolve team ID for embedding role resolution (pick first team)
	var teamID *string
	if len(teamIDs) > 0 {
		teamID = &teamIDs[0]
	}

	// Embed the query
	embedResult, err := t.embedder.EmbedChunks(ctx, execCtx.UserID, teamID, []string{args.Query})
	if err != nil {
		return "", fmt.Errorf("failed to embed search query: %w", err)
	}
	if len(embedResult.Vectors) == 0 {
		return "", fmt.Errorf("embedding produced no vectors")
	}

	queryVec := embedResult.Vectors[0]

	// Track embedding usage (role = "embedding")
	chanPtr := &execCtx.ChannelID
	t.embedder.LogUsage(ctx, execCtx.UserID, chanPtr, embedResult)

	// Similarity search across all active KBs
	const defaultThreshold = 0.3
	results, err := t.stores.KnowledgeBases.SimilaritySearch(ctx, kbIDs, queryVec, defaultThreshold, args.MaxResults)
	if err != nil {
		return "", fmt.Errorf("search failed: %w", err)
	}

	// Collect searched KB names for attribution
	searchedKBs := make(map[string]bool)
	for _, r := range results {
		searchedKBs[r.KBName] = true
	}
	kbNames := make([]string, 0, len(searchedKBs))
	for name := range searchedKBs {
		kbNames = append(kbNames, name)
	}

	// If no results found from search, list all searched KB names
	if len(kbNames) == 0 {
		for _, kbID := range kbIDs {
			kb, err := t.stores.KnowledgeBases.GetByID(ctx, kbID)
			if err == nil {
				kbNames = append(kbNames, kb.Name)
			}
		}
	}

	log.Printf("🔍 kb_search: %q → %d results across %d KBs", args.Query, len(results), len(kbIDs))

	out, _ := json.Marshal(map[string]interface{}{
		"results":      results,
		"query":        args.Query,
		"total":        len(results),
		"searched_kbs": kbNames,
	})
	return string(out), nil
}