Changeset 0.14.0 (#67)

server/tools/kbsearch.go

@@ -0,0 +1,159 @@
+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 (respects scope: global, team, personal)
+	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 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
+}