Changeset 0.14.0 (#67)
This commit is contained in:
99
server/database/migrations/008_knowledge_bases.sql
Normal file
99
server/database/migrations/008_knowledge_bases.sql
Normal file
@@ -0,0 +1,99 @@
|
||||
-- 008_knowledge_bases.sql
|
||||
-- Knowledge Bases: KB metadata, documents, chunks, channel links.
|
||||
-- Requires: pgvector extension already installed (via db-bootstrap.sh).
|
||||
|
||||
-- ── Knowledge Bases ──────────────────────────
|
||||
|
||||
CREATE TABLE IF NOT EXISTS knowledge_bases (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
name TEXT NOT NULL,
|
||||
description TEXT NOT NULL DEFAULT '',
|
||||
scope TEXT NOT NULL DEFAULT 'global', -- global, team, personal
|
||||
owner_id UUID REFERENCES users(id) ON DELETE CASCADE,
|
||||
team_id UUID REFERENCES teams(id) ON DELETE CASCADE,
|
||||
|
||||
-- Snapshot of embedding config at creation time.
|
||||
-- Changing the model requires a full re-embed (rebuild endpoint).
|
||||
-- { "provider_config_id": "...", "model_id": "...", "dimensions": 1536 }
|
||||
embedding_config JSONB NOT NULL DEFAULT '{}',
|
||||
|
||||
-- Denormalized stats, updated by ingest pipeline.
|
||||
document_count INT NOT NULL DEFAULT 0,
|
||||
chunk_count INT NOT NULL DEFAULT 0,
|
||||
total_bytes BIGINT NOT NULL DEFAULT 0,
|
||||
|
||||
status TEXT NOT NULL DEFAULT 'active', -- active, processing, error
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
|
||||
CONSTRAINT kb_scope_check CHECK (
|
||||
(scope = 'global' AND owner_id IS NULL) OR
|
||||
(scope = 'team' AND team_id IS NOT NULL) OR
|
||||
(scope = 'personal' AND owner_id IS NOT NULL)
|
||||
)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_kb_scope ON knowledge_bases(scope);
|
||||
CREATE INDEX IF NOT EXISTS idx_kb_owner ON knowledge_bases(owner_id) WHERE owner_id IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_kb_team ON knowledge_bases(team_id) WHERE team_id IS NOT NULL;
|
||||
|
||||
-- ── KB Documents ─────────────────────────────
|
||||
-- One row per uploaded file. Blobs live in ObjectStore under:
|
||||
-- kb/{kb_id}/{doc_id}_{filename}
|
||||
|
||||
CREATE TABLE IF NOT EXISTS kb_documents (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
kb_id UUID NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
||||
filename TEXT NOT NULL,
|
||||
content_type TEXT NOT NULL,
|
||||
size_bytes BIGINT NOT NULL,
|
||||
storage_key TEXT NOT NULL,
|
||||
extracted_text TEXT, -- full text for re-chunking
|
||||
chunk_count INT NOT NULL DEFAULT 0,
|
||||
status TEXT NOT NULL DEFAULT 'pending', -- pending, chunking, embedding, ready, error
|
||||
error TEXT,
|
||||
uploaded_by UUID NOT NULL REFERENCES users(id),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_kbdoc_kb ON kb_documents(kb_id);
|
||||
|
||||
-- ── KB Chunks ────────────────────────────────
|
||||
-- Chunked text + embedding vector for similarity search.
|
||||
-- vector(3072) accommodates all current embedding models:
|
||||
-- OpenAI ada-002 / text-embedding-3-small = 1536
|
||||
-- OpenAI text-embedding-3-large = 3072
|
||||
-- Open-source models = typically 768-1024
|
||||
-- Smaller vectors are zero-padded at insert time.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS kb_chunks (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
kb_id UUID NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
||||
document_id UUID NOT NULL REFERENCES kb_documents(id) ON DELETE CASCADE,
|
||||
chunk_index INT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
token_count INT NOT NULL DEFAULT 0,
|
||||
embedding VECTOR(3072),
|
||||
metadata JSONB NOT NULL DEFAULT '{}',
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_kbchunk_kb ON kb_chunks(kb_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_kbchunk_doc ON kb_chunks(document_id);
|
||||
|
||||
-- NOTE: IVFFlat index for similarity search is created after first data
|
||||
-- load (needs rows to train). The ingest handler creates it:
|
||||
-- CREATE INDEX idx_kbchunk_embedding ON kb_chunks
|
||||
-- USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
|
||||
|
||||
-- ── Channel-KB Links ─────────────────────────
|
||||
-- Which KBs are active for a given channel.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS channel_knowledge_bases (
|
||||
channel_id UUID NOT NULL REFERENCES channels(id) ON DELETE CASCADE,
|
||||
kb_id UUID NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
||||
enabled BOOLEAN NOT NULL DEFAULT true,
|
||||
added_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (channel_id, kb_id)
|
||||
);
|
||||
12
server/database/migrations/009_notes_embedding.sql
Normal file
12
server/database/migrations/009_notes_embedding.sql
Normal file
@@ -0,0 +1,12 @@
|
||||
-- 009_notes_embedding.sql
|
||||
-- Adds vector embedding column to notes for semantic search.
|
||||
-- Uses the same vector(3072) dimension as kb_chunks for uniformity.
|
||||
|
||||
ALTER TABLE notes ADD COLUMN IF NOT EXISTS embedding vector(3072);
|
||||
|
||||
-- NOTE: pgvector HNSW indexes are limited to 2000 dimensions.
|
||||
-- For 3072-dim vectors, IVFFlat is the option but requires training
|
||||
-- data (rows). The notes table is typically small enough that a
|
||||
-- sequential scan is fast. If needed, create after data exists:
|
||||
-- CREATE INDEX idx_notes_embedding ON notes
|
||||
-- USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
|
||||
@@ -62,16 +62,23 @@ func SetupTestDB() func() {
|
||||
return func() {}
|
||||
}
|
||||
|
||||
// Attempt to create the test DB. If the user lacks CREATE DATABASE
|
||||
// permissions (e.g. CI already created it via admin bootstrap step),
|
||||
// that's fine — we'll just connect to the existing one.
|
||||
// Check if test DB already exists (e.g. created by CI bootstrap with
|
||||
// admin privileges and extensions like pgvector already installed).
|
||||
// Only create if it doesn't exist — never drop, because recreating
|
||||
// as the app user would lose extensions that require superuser.
|
||||
createdByUs := false
|
||||
adminDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", testDBName))
|
||||
if _, err := adminDB.Exec(fmt.Sprintf("CREATE DATABASE %s", testDBName)); err != nil {
|
||||
log.Printf("⚠ Cannot CREATE DATABASE %s (will try to connect to existing): %v", testDBName, err)
|
||||
var dbExists bool
|
||||
err = adminDB.QueryRow("SELECT 1 FROM pg_database WHERE datname = $1", testDBName).Scan(&dbExists)
|
||||
if err != nil {
|
||||
// DB doesn't exist — create it
|
||||
if _, err := adminDB.Exec(fmt.Sprintf("CREATE DATABASE %s", testDBName)); err != nil {
|
||||
log.Printf("⚠ Cannot CREATE DATABASE %s (will try to connect to existing): %v", testDBName, err)
|
||||
} else {
|
||||
createdByUs = true
|
||||
log.Printf("✓ Created test database: %s", testDBName)
|
||||
}
|
||||
} else {
|
||||
createdByUs = true
|
||||
log.Printf("✓ Created test database: %s", testDBName)
|
||||
log.Printf("✓ Test database %s already exists (keeping extensions)", testDBName)
|
||||
}
|
||||
|
||||
// Build DSN for the test DB
|
||||
@@ -100,7 +107,7 @@ func SetupTestDB() func() {
|
||||
TestDB = DB
|
||||
|
||||
// Ensure required extensions (may already exist from CI bootstrap)
|
||||
for _, ext := range []string{"uuid-ossp", "pgcrypto"} {
|
||||
for _, ext := range []string{"uuid-ossp", "pgcrypto", "vector"} {
|
||||
DB.Exec(fmt.Sprintf(`CREATE EXTENSION IF NOT EXISTS "%s"`, ext))
|
||||
}
|
||||
|
||||
@@ -166,6 +173,10 @@ func TruncateAll(t *testing.T) {
|
||||
"messages",
|
||||
"channel_models",
|
||||
"channel_members",
|
||||
"channel_knowledge_bases",
|
||||
"kb_chunks",
|
||||
"kb_documents",
|
||||
"knowledge_bases",
|
||||
"channels",
|
||||
"user_model_settings",
|
||||
"model_catalog",
|
||||
|
||||
@@ -780,6 +780,14 @@ func (h *CompletionHandler) loadConversation(channelID, userID, presetSystemProm
|
||||
})
|
||||
}
|
||||
|
||||
// ── KB hint (nudge LLM to use kb_search when KBs are active) ──
|
||||
if kbHint := BuildKBHint(context.Background(), h.stores, channelID, userID); 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 {
|
||||
|
||||
@@ -135,6 +135,20 @@ func setupHarness(t *testing.T) *testHarness {
|
||||
protected.GET("/channels/:id", channels.GetChannel)
|
||||
protected.DELETE("/channels/:id", channels.DeleteChannel)
|
||||
|
||||
// Knowledge Bases (nil storage/ingester/embedder — CRUD works, upload/rebuild return 503)
|
||||
kbH := NewKnowledgeBaseHandler(stores, nil, nil, nil)
|
||||
protected.POST("/knowledge-bases", kbH.CreateKB)
|
||||
protected.GET("/knowledge-bases", kbH.ListKBs)
|
||||
protected.GET("/knowledge-bases/:id", kbH.GetKB)
|
||||
protected.PUT("/knowledge-bases/:id", kbH.UpdateKB)
|
||||
protected.DELETE("/knowledge-bases/:id", kbH.DeleteKB)
|
||||
protected.POST("/knowledge-bases/:id/documents", kbH.UploadDocument)
|
||||
protected.GET("/knowledge-bases/:id/documents", kbH.ListDocuments)
|
||||
protected.POST("/knowledge-bases/:id/search", kbH.SearchKB)
|
||||
protected.POST("/knowledge-bases/:id/rebuild", kbH.RebuildKB)
|
||||
protected.GET("/channels/:id/knowledge-bases", kbH.GetChannelKBs)
|
||||
protected.PUT("/channels/:id/knowledge-bases", kbH.SetChannelKBs)
|
||||
|
||||
// Attachments (nil storage = upload returns 503, but metadata works)
|
||||
attachH := NewAttachmentHandler(stores, nil, nil)
|
||||
protected.POST("/channels/:id/attachments", attachH.Upload)
|
||||
@@ -2230,3 +2244,202 @@ func TestIntegration_Roles_GetSingleRole(t *testing.T) {
|
||||
t.Fatalf("get seeded-but-empty role: want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// ── KB CRUD ─────────────────────────────────────
|
||||
|
||||
func TestIntegration_KnowledgeBaseCRUD(t *testing.T) {
|
||||
h := setupHarness(t)
|
||||
_, adminToken := h.createAdminUser("admin", "admin@test.com")
|
||||
|
||||
// Create personal KB
|
||||
w := h.request("POST", "/api/v1/knowledge-bases", adminToken, map[string]interface{}{
|
||||
"name": "Test KB", "description": "A test knowledge base",
|
||||
})
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("create KB: want 201, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var kb map[string]interface{}
|
||||
decode(w, &kb)
|
||||
kbID := kb["id"].(string)
|
||||
|
||||
if kb["name"] != "Test KB" {
|
||||
t.Fatalf("create KB: name mismatch: got %s", kb["name"])
|
||||
}
|
||||
if kb["scope"] != "personal" {
|
||||
t.Fatalf("create KB: scope should default to personal: got %s", kb["scope"])
|
||||
}
|
||||
|
||||
// List KBs
|
||||
w = h.request("GET", "/api/v1/knowledge-bases", adminToken, nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("list KBs: want 200, got %d", w.Code)
|
||||
}
|
||||
var kbList map[string]interface{}
|
||||
decode(w, &kbList)
|
||||
items := kbList["data"].([]interface{})
|
||||
if len(items) != 1 {
|
||||
t.Fatalf("list KBs: expected 1, got %d", len(items))
|
||||
}
|
||||
|
||||
// Get single KB
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/knowledge-bases/%s", kbID), adminToken, nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("get KB: want 200, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Update KB
|
||||
w = h.request("PUT", fmt.Sprintf("/api/v1/knowledge-bases/%s", kbID), adminToken,
|
||||
map[string]interface{}{"name": "Updated KB"})
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("update KB: want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// Verify update
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/knowledge-bases/%s", kbID), adminToken, nil)
|
||||
var updated map[string]interface{}
|
||||
decode(w, &updated)
|
||||
if updated["name"] != "Updated KB" {
|
||||
t.Fatalf("update KB: name not updated, got %s", updated["name"])
|
||||
}
|
||||
|
||||
// Upload document — should 503 (no object store in test harness)
|
||||
w = h.request("POST", fmt.Sprintf("/api/v1/knowledge-bases/%s/documents", kbID), adminToken, nil)
|
||||
if w.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("upload doc without objstore: want 503, got %d", w.Code)
|
||||
}
|
||||
|
||||
// List documents (empty)
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/knowledge-bases/%s/documents", kbID), adminToken, nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("list docs: want 200, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Rebuild — should 503 (no ingester in test harness)
|
||||
w = h.request("POST", fmt.Sprintf("/api/v1/knowledge-bases/%s/rebuild", kbID), adminToken, nil)
|
||||
if w.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("rebuild without ingester: want 503, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Delete KB
|
||||
w = h.request("DELETE", fmt.Sprintf("/api/v1/knowledge-bases/%s", kbID), adminToken, nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("delete KB: want 200, got %d", w.Code)
|
||||
}
|
||||
|
||||
// Verify deleted
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/knowledge-bases/%s", kbID), adminToken, nil)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("get deleted KB: want 404, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// ── KB Cross-User Isolation ─────────────────────
|
||||
|
||||
func TestIntegration_KBCrossUserIsolation(t *testing.T) {
|
||||
h := setupHarness(t)
|
||||
_, adminToken := h.createAdminUser("admin", "admin@test.com")
|
||||
|
||||
database.TestDB.Exec("INSERT INTO platform_policies (key, value) VALUES ('allow_registration', 'true') ON CONFLICT (key) DO UPDATE SET value = 'true'")
|
||||
database.TestDB.Exec("INSERT INTO platform_policies (key, value) VALUES ('default_user_active', 'true') ON CONFLICT (key) DO UPDATE SET value = 'true'")
|
||||
|
||||
_, userToken := h.registerUser("alice", "alice@test.com", "password123")
|
||||
|
||||
// Admin creates personal KB
|
||||
w := h.request("POST", "/api/v1/knowledge-bases", adminToken, map[string]interface{}{
|
||||
"name": "Admin Private KB",
|
||||
})
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("admin create KB: want 201, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var adminKB map[string]interface{}
|
||||
decode(w, &adminKB)
|
||||
adminKBID := adminKB["id"].(string)
|
||||
|
||||
// Alice should NOT see admin's personal KB
|
||||
w = h.request("GET", "/api/v1/knowledge-bases", userToken, nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("alice list KBs: want 200, got %d", w.Code)
|
||||
}
|
||||
var aliceList map[string]interface{}
|
||||
decode(w, &aliceList)
|
||||
items := aliceList["data"].([]interface{})
|
||||
if len(items) != 0 {
|
||||
t.Fatalf("alice should see 0 KBs, got %d", len(items))
|
||||
}
|
||||
|
||||
// Alice should NOT be able to access admin's KB directly
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/knowledge-bases/%s", adminKBID), userToken, nil)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("alice access admin KB: want 404 (hidden), got %d", w.Code)
|
||||
}
|
||||
|
||||
// Alice should NOT be able to delete admin's KB
|
||||
w = h.request("DELETE", fmt.Sprintf("/api/v1/knowledge-bases/%s", adminKBID), userToken, nil)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("alice delete admin KB: want 404 (hidden), got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Channel KB Linking ──────────────────────────
|
||||
|
||||
func TestIntegration_ChannelKBLinking(t *testing.T) {
|
||||
h := setupHarness(t)
|
||||
_, adminToken := h.createAdminUser("admin", "admin@test.com")
|
||||
|
||||
// Create a channel
|
||||
w := h.request("POST", "/api/v1/channels", adminToken, map[string]interface{}{
|
||||
"title": "Test Channel",
|
||||
})
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("create channel: want 201, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var ch map[string]interface{}
|
||||
decode(w, &ch)
|
||||
channelID := ch["id"].(string)
|
||||
|
||||
// Create a KB
|
||||
w = h.request("POST", "/api/v1/knowledge-bases", adminToken, map[string]interface{}{
|
||||
"name": "Channel KB",
|
||||
})
|
||||
if w.Code != http.StatusCreated {
|
||||
t.Fatalf("create KB: want 201, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var kb map[string]interface{}
|
||||
decode(w, &kb)
|
||||
kbID := kb["id"].(string)
|
||||
|
||||
// Link KB to channel
|
||||
w = h.request("PUT", fmt.Sprintf("/api/v1/channels/%s/knowledge-bases", channelID), adminToken,
|
||||
map[string]interface{}{"kb_ids": []string{kbID}})
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("link KB to channel: want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// Get channel KBs
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/channels/%s/knowledge-bases", channelID), adminToken, nil)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("get channel KBs: want 200, got %d", w.Code)
|
||||
}
|
||||
var linked map[string]interface{}
|
||||
decode(w, &linked)
|
||||
data := linked["data"].([]interface{})
|
||||
if len(data) != 1 {
|
||||
t.Fatalf("channel should have 1 linked KB, got %d", len(data))
|
||||
}
|
||||
|
||||
// Unlink (set empty)
|
||||
w = h.request("PUT", fmt.Sprintf("/api/v1/channels/%s/knowledge-bases", channelID), adminToken,
|
||||
map[string]interface{}{"kb_ids": []string{}})
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("unlink KBs: want 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// Verify empty
|
||||
w = h.request("GET", fmt.Sprintf("/api/v1/channels/%s/knowledge-bases", channelID), adminToken, nil)
|
||||
var empty map[string]interface{}
|
||||
decode(w, &empty)
|
||||
emptyData := empty["data"].([]interface{})
|
||||
if len(emptyData) != 0 {
|
||||
t.Fatalf("channel should have 0 linked KBs after unlink, got %d", len(emptyData))
|
||||
}
|
||||
}
|
||||
|
||||
950
server/handlers/knowledge_bases.go
Normal file
950
server/handlers/knowledge_bases.go
Normal file
@@ -0,0 +1,950 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/database"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/knowledge"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/models"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/storage"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/store"
|
||||
)
|
||||
|
||||
// ── Limits ───────────────────────────────────
|
||||
|
||||
const (
|
||||
kbMaxDocSize = 10 * 1024 * 1024 // 10 MB per document
|
||||
kbMaxNameLen = 200
|
||||
kbMaxDescLen = 2000
|
||||
)
|
||||
|
||||
// allowedKBMIMETypes lists content types accepted for KB documents.
|
||||
var allowedKBMIMETypes = map[string]bool{
|
||||
"text/plain": true,
|
||||
"text/markdown": true,
|
||||
"text/csv": true,
|
||||
"text/html": true,
|
||||
// Future: PDF, docx once extraction sidecar is wired in
|
||||
}
|
||||
|
||||
// ── Request / Response Types ─────────────────
|
||||
|
||||
type createKBRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Scope string `json:"scope"` // personal (default), team, global
|
||||
TeamID string `json:"team_id"`
|
||||
}
|
||||
|
||||
type updateKBRequest struct {
|
||||
Name *string `json:"name"`
|
||||
Description *string `json:"description"`
|
||||
}
|
||||
|
||||
type kbResponse struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Scope string `json:"scope"`
|
||||
OwnerID *string `json:"owner_id,omitempty"`
|
||||
TeamID *string `json:"team_id,omitempty"`
|
||||
EmbeddingConfig models.JSONMap `json:"embedding_config"`
|
||||
DocumentCount int `json:"document_count"`
|
||||
ChunkCount int `json:"chunk_count"`
|
||||
TotalBytes int64 `json:"total_bytes"`
|
||||
Status string `json:"status"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
type docResponse struct {
|
||||
ID string `json:"id"`
|
||||
KBID string `json:"kb_id"`
|
||||
Filename string `json:"filename"`
|
||||
ContentType string `json:"content_type"`
|
||||
SizeBytes int64 `json:"size_bytes"`
|
||||
ChunkCount int `json:"chunk_count"`
|
||||
Status string `json:"status"`
|
||||
Error *string `json:"error,omitempty"`
|
||||
UploadedBy string `json:"uploaded_by"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
type searchKBRequest struct {
|
||||
Query string `json:"query" binding:"required"`
|
||||
Limit int `json:"limit"`
|
||||
Threshold float64 `json:"threshold"`
|
||||
}
|
||||
|
||||
// ── Handler ──────────────────────────────────
|
||||
|
||||
// KnowledgeBaseHandler handles KB CRUD and document management.
|
||||
type KnowledgeBaseHandler struct {
|
||||
stores store.Stores
|
||||
objStore storage.ObjectStore
|
||||
ingester *knowledge.Ingester
|
||||
embedder *knowledge.Embedder
|
||||
}
|
||||
|
||||
// NewKnowledgeBaseHandler creates a KB handler.
|
||||
func NewKnowledgeBaseHandler(
|
||||
stores store.Stores,
|
||||
objStore storage.ObjectStore,
|
||||
ingester *knowledge.Ingester,
|
||||
embedder *knowledge.Embedder,
|
||||
) *KnowledgeBaseHandler {
|
||||
return &KnowledgeBaseHandler{
|
||||
stores: stores,
|
||||
objStore: objStore,
|
||||
ingester: ingester,
|
||||
embedder: embedder,
|
||||
}
|
||||
}
|
||||
|
||||
// ── KB CRUD ──────────────────────────────────
|
||||
|
||||
// CreateKB creates a new knowledge base.
|
||||
// POST /api/v1/knowledge-bases
|
||||
func (h *KnowledgeBaseHandler) CreateKB(c *gin.Context) {
|
||||
var req createKBRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if len(req.Name) > kbMaxNameLen {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name too long"})
|
||||
return
|
||||
}
|
||||
if len(req.Description) > kbMaxDescLen {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "description too long"})
|
||||
return
|
||||
}
|
||||
|
||||
userID := getUserID(c)
|
||||
|
||||
// Default scope
|
||||
scope := "personal"
|
||||
if req.Scope != "" {
|
||||
scope = req.Scope
|
||||
}
|
||||
|
||||
// Validate scope
|
||||
switch scope {
|
||||
case "personal":
|
||||
// OK — owned by user
|
||||
case "team":
|
||||
if req.TeamID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "team_id required for team scope"})
|
||||
return
|
||||
}
|
||||
// TODO: verify user is team admin
|
||||
case "global":
|
||||
// TODO: verify user is admin
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "scope must be personal, team, or global"})
|
||||
return
|
||||
}
|
||||
|
||||
kb := &models.KnowledgeBase{
|
||||
Name: req.Name,
|
||||
Description: req.Description,
|
||||
Scope: scope,
|
||||
EmbeddingConfig: models.JSONMap{},
|
||||
Status: "active",
|
||||
}
|
||||
|
||||
// Set ownership based on scope
|
||||
switch scope {
|
||||
case "personal":
|
||||
kb.OwnerID = &userID
|
||||
case "team":
|
||||
kb.TeamID = &req.TeamID
|
||||
}
|
||||
|
||||
if err := h.stores.KnowledgeBases.Create(c.Request.Context(), kb); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create knowledge base"})
|
||||
return
|
||||
}
|
||||
|
||||
h.auditLog(c, "kb.create", kb.ID, map[string]interface{}{
|
||||
"name": kb.Name, "scope": kb.Scope,
|
||||
})
|
||||
|
||||
c.JSON(http.StatusCreated, toKBResponse(kb))
|
||||
}
|
||||
|
||||
// ListKBs lists knowledge bases accessible to the current user.
|
||||
// GET /api/v1/knowledge-bases
|
||||
func (h *KnowledgeBaseHandler) ListKBs(c *gin.Context) {
|
||||
userID := getUserID(c)
|
||||
|
||||
// Get user's team IDs for scoped listing
|
||||
teamIDs := h.getUserTeamIDs(c, userID)
|
||||
|
||||
kbs, err := h.stores.KnowledgeBases.ListForUser(c.Request.Context(), userID, teamIDs)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list knowledge bases"})
|
||||
return
|
||||
}
|
||||
|
||||
items := make([]kbResponse, len(kbs))
|
||||
for i := range kbs {
|
||||
items[i] = toKBResponse(&kbs[i])
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"data": items, "total": len(items)})
|
||||
}
|
||||
|
||||
// GetKB returns a single knowledge base.
|
||||
// GET /api/v1/knowledge-bases/:id
|
||||
func (h *KnowledgeBaseHandler) GetKB(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, toKBResponse(kb))
|
||||
}
|
||||
|
||||
// UpdateKB updates a knowledge base's name/description.
|
||||
// PUT /api/v1/knowledge-bases/:id
|
||||
func (h *KnowledgeBaseHandler) UpdateKB(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req updateKBRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
fields := make(map[string]interface{})
|
||||
if req.Name != nil {
|
||||
if len(*req.Name) > kbMaxNameLen {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "name too long"})
|
||||
return
|
||||
}
|
||||
fields["name"] = *req.Name
|
||||
}
|
||||
if req.Description != nil {
|
||||
if len(*req.Description) > kbMaxDescLen {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "description too long"})
|
||||
return
|
||||
}
|
||||
fields["description"] = *req.Description
|
||||
}
|
||||
|
||||
if len(fields) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "no fields to update"})
|
||||
return
|
||||
}
|
||||
|
||||
if err := h.stores.KnowledgeBases.Update(c.Request.Context(), kb.ID, fields); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update knowledge base"})
|
||||
return
|
||||
}
|
||||
|
||||
// Reload for response
|
||||
updated, err := h.stores.KnowledgeBases.GetByID(c.Request.Context(), kb.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to reload knowledge base"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, toKBResponse(updated))
|
||||
}
|
||||
|
||||
// DeleteKB deletes a knowledge base and all its documents/chunks.
|
||||
// DELETE /api/v1/knowledge-bases/:id
|
||||
func (h *KnowledgeBaseHandler) DeleteKB(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
// Delete stored files
|
||||
if h.objStore != nil {
|
||||
prefix := fmt.Sprintf("kb/%s/", kb.ID)
|
||||
_ = h.objStore.DeletePrefix(c.Request.Context(), prefix)
|
||||
}
|
||||
|
||||
// Delete from DB (cascades to documents and chunks)
|
||||
if err := h.stores.KnowledgeBases.Delete(c.Request.Context(), kb.ID); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete knowledge base"})
|
||||
return
|
||||
}
|
||||
|
||||
h.auditLog(c, "kb.delete", kb.ID, map[string]interface{}{
|
||||
"name": kb.Name, "scope": kb.Scope,
|
||||
})
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
// ── Document Management ──────────────────────
|
||||
|
||||
// UploadDocument uploads a document to a KB and starts async ingestion.
|
||||
// POST /api/v1/knowledge-bases/:id/documents
|
||||
func (h *KnowledgeBaseHandler) UploadDocument(c *gin.Context) {
|
||||
if h.objStore == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "file storage not configured"})
|
||||
return
|
||||
}
|
||||
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
// Verify embedding role is configured
|
||||
if !h.embedder.IsConfigured(c.Request.Context()) {
|
||||
c.JSON(http.StatusPreconditionFailed, gin.H{
|
||||
"error": "embedding model role not configured — set up an embedding provider in Settings → Model Roles",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Parse multipart
|
||||
file, header, err := c.Request.FormFile("file")
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing file field"})
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
if header.Size > kbMaxDocSize {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"error": fmt.Sprintf("file too large (max %d MB)", kbMaxDocSize/(1024*1024)),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// MIME detection
|
||||
buf := make([]byte, 512)
|
||||
n, _ := file.Read(buf)
|
||||
contentType := http.DetectContentType(buf[:n])
|
||||
if idx := strings.Index(contentType, ";"); idx > 0 {
|
||||
contentType = strings.TrimSpace(contentType[:idx])
|
||||
}
|
||||
|
||||
// Fall back to extension for octet-stream
|
||||
if contentType == "application/octet-stream" {
|
||||
ext := strings.ToLower(filepath.Ext(header.Filename))
|
||||
switch ext {
|
||||
case ".md", ".markdown":
|
||||
contentType = "text/markdown"
|
||||
case ".csv":
|
||||
contentType = "text/csv"
|
||||
case ".html", ".htm":
|
||||
contentType = "text/html"
|
||||
case ".txt":
|
||||
contentType = "text/plain"
|
||||
}
|
||||
}
|
||||
|
||||
// Allowlist check
|
||||
if !allowedKBMIMETypes[contentType] {
|
||||
c.JSON(http.StatusBadRequest, gin.H{
|
||||
"error": fmt.Sprintf("file type %q not supported for knowledge bases — only text, markdown, CSV, and HTML are currently supported", contentType),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Reset reader
|
||||
if _, err := file.Seek(0, io.SeekStart); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to process file"})
|
||||
return
|
||||
}
|
||||
|
||||
userID := getUserID(c)
|
||||
filename := sanitizeFilename(header.Filename)
|
||||
|
||||
// Create document record (status=pending)
|
||||
doc := &models.KBDocument{
|
||||
KBID: kb.ID,
|
||||
Filename: filename,
|
||||
ContentType: contentType,
|
||||
SizeBytes: header.Size,
|
||||
StorageKey: "placeholder",
|
||||
Status: "pending",
|
||||
UploadedBy: userID,
|
||||
}
|
||||
|
||||
if err := h.stores.KnowledgeBases.CreateDocument(c.Request.Context(), doc); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create document record"})
|
||||
return
|
||||
}
|
||||
|
||||
// Build storage key and save file
|
||||
doc.StorageKey = fmt.Sprintf("kb/%s/%s_%s", kb.ID, doc.ID, filename)
|
||||
if err := h.objStore.Put(c.Request.Context(), doc.StorageKey, file, header.Size, contentType); err != nil {
|
||||
// Rollback DB record
|
||||
h.stores.KnowledgeBases.DeleteDocument(c.Request.Context(), doc.ID)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to store file"})
|
||||
return
|
||||
}
|
||||
|
||||
// Update storage key in DB
|
||||
h.updateDocStorageKey(c, doc.ID, doc.StorageKey)
|
||||
|
||||
// Resolve team ID for embedding role
|
||||
var teamID *string
|
||||
if kb.TeamID != nil {
|
||||
teamID = kb.TeamID
|
||||
}
|
||||
|
||||
// Fire async ingestion
|
||||
h.ingester.IngestDocument(kb, doc, userID, teamID)
|
||||
|
||||
h.auditLog(c, "kb.document.upload", kb.ID, map[string]interface{}{
|
||||
"kb_name": kb.Name, "doc_id": doc.ID, "filename": doc.Filename,
|
||||
"size_bytes": doc.SizeBytes,
|
||||
})
|
||||
|
||||
c.JSON(http.StatusAccepted, toDocResponse(doc))
|
||||
}
|
||||
|
||||
// ListDocuments lists all documents in a KB.
|
||||
// GET /api/v1/knowledge-bases/:id/documents
|
||||
func (h *KnowledgeBaseHandler) ListDocuments(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
docs, err := h.stores.KnowledgeBases.ListDocuments(c.Request.Context(), kb.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list documents"})
|
||||
return
|
||||
}
|
||||
|
||||
items := make([]docResponse, len(docs))
|
||||
for i := range docs {
|
||||
items[i] = toDocResponse(&docs[i])
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"data": items, "total": len(items)})
|
||||
}
|
||||
|
||||
// GetDocumentStatus returns the current ingestion status of a document.
|
||||
// GET /api/v1/knowledge-bases/:id/documents/:docId/status
|
||||
func (h *KnowledgeBaseHandler) GetDocumentStatus(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
docID := c.Param("docId")
|
||||
doc, err := h.stores.KnowledgeBases.GetDocument(c.Request.Context(), docID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "document not found"})
|
||||
return
|
||||
}
|
||||
|
||||
if doc.KBID != kb.ID {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "document not found"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"id": doc.ID,
|
||||
"status": doc.Status,
|
||||
"chunk_count": doc.ChunkCount,
|
||||
"error": doc.Error,
|
||||
"filename": doc.Filename,
|
||||
})
|
||||
}
|
||||
|
||||
// DeleteDocument removes a document and its chunks.
|
||||
// DELETE /api/v1/knowledge-bases/:id/documents/:docId
|
||||
func (h *KnowledgeBaseHandler) DeleteDocument(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
docID := c.Param("docId")
|
||||
doc, err := h.stores.KnowledgeBases.GetDocument(c.Request.Context(), docID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "document not found"})
|
||||
return
|
||||
}
|
||||
|
||||
if doc.KBID != kb.ID {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "document not found"})
|
||||
return
|
||||
}
|
||||
|
||||
// Delete chunks first
|
||||
h.stores.KnowledgeBases.DeleteChunksForDocument(c.Request.Context(), docID)
|
||||
|
||||
// Delete the document record (returns doc for storage cleanup)
|
||||
deleted, err := h.stores.KnowledgeBases.DeleteDocument(c.Request.Context(), docID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to delete document"})
|
||||
return
|
||||
}
|
||||
|
||||
// Clean up stored file
|
||||
if h.objStore != nil && deleted != nil && deleted.StorageKey != "" && deleted.StorageKey != "placeholder" {
|
||||
_ = h.objStore.Delete(c.Request.Context(), deleted.StorageKey)
|
||||
}
|
||||
|
||||
// Update KB stats
|
||||
h.stores.KnowledgeBases.UpdateStats(c.Request.Context(), kb.ID)
|
||||
|
||||
h.auditLog(c, "kb.document.delete", kb.ID, map[string]interface{}{
|
||||
"kb_name": kb.Name, "doc_id": docID,
|
||||
"filename": safeFilename(deleted),
|
||||
})
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"deleted": true})
|
||||
}
|
||||
|
||||
// ── Search ───────────────────────────────────
|
||||
|
||||
// SearchKB performs a direct similarity search against a KB.
|
||||
// POST /api/v1/knowledge-bases/:id/search
|
||||
func (h *KnowledgeBaseHandler) SearchKB(c *gin.Context) {
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
var req searchKBRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Limit <= 0 || req.Limit > 20 {
|
||||
req.Limit = 5
|
||||
}
|
||||
if req.Threshold <= 0 {
|
||||
req.Threshold = 0.3 // default cosine similarity threshold
|
||||
}
|
||||
|
||||
userID := getUserID(c)
|
||||
var teamID *string
|
||||
if kb.TeamID != nil {
|
||||
teamID = kb.TeamID
|
||||
}
|
||||
|
||||
// Embed the query
|
||||
embedResult, err := h.embedder.EmbedChunks(c.Request.Context(), userID, teamID, []string{req.Query})
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to embed query"})
|
||||
return
|
||||
}
|
||||
|
||||
if len(embedResult.Vectors) == 0 {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "embedding produced no vectors"})
|
||||
return
|
||||
}
|
||||
|
||||
// Search
|
||||
results, err := h.stores.KnowledgeBases.SimilaritySearch(
|
||||
c.Request.Context(),
|
||||
[]string{kb.ID},
|
||||
embedResult.Vectors[0],
|
||||
req.Threshold,
|
||||
req.Limit,
|
||||
)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "search failed"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"data": results,
|
||||
"query": req.Query,
|
||||
"total": len(results),
|
||||
})
|
||||
}
|
||||
|
||||
// ── Helpers ──────────────────────────────────
|
||||
|
||||
// loadAndAuthorize loads a KB by ID from the URL and checks the user
|
||||
// has access to it (owner, team member, or admin for global).
|
||||
func (h *KnowledgeBaseHandler) loadAndAuthorize(c *gin.Context) (*models.KnowledgeBase, bool) {
|
||||
kbID := c.Param("id")
|
||||
if kbID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing knowledge base ID"})
|
||||
return nil, false
|
||||
}
|
||||
|
||||
kb, err := h.stores.KnowledgeBases.GetByID(c.Request.Context(), kbID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "knowledge base not found"})
|
||||
return nil, false
|
||||
}
|
||||
|
||||
userID := getUserID(c)
|
||||
|
||||
switch kb.Scope {
|
||||
case "personal":
|
||||
if kb.OwnerID == nil || *kb.OwnerID != userID {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "knowledge base not found"})
|
||||
return nil, false
|
||||
}
|
||||
case "team":
|
||||
// Verify user is a member of the team
|
||||
teamIDs := h.getUserTeamIDs(c, userID)
|
||||
found := false
|
||||
for _, tid := range teamIDs {
|
||||
if kb.TeamID != nil && tid == *kb.TeamID {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "knowledge base not found"})
|
||||
return nil, false
|
||||
}
|
||||
case "global":
|
||||
// Global KBs are readable by everyone, writable by admins
|
||||
// For now, allow all authenticated users to read
|
||||
}
|
||||
|
||||
return kb, true
|
||||
}
|
||||
|
||||
// getUserTeamIDs returns the team IDs for the current user.
|
||||
func (h *KnowledgeBaseHandler) getUserTeamIDs(c *gin.Context, userID string) []string {
|
||||
ids, err := h.stores.Teams.GetUserTeamIDs(c.Request.Context(), userID)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// updateDocStorageKey sets the storage_key on a document after upload.
|
||||
func (h *KnowledgeBaseHandler) updateDocStorageKey(c *gin.Context, docID, storageKey string) {
|
||||
database.DB.ExecContext(c.Request.Context(),
|
||||
`UPDATE kb_documents SET storage_key = $1 WHERE id = $2`,
|
||||
storageKey, docID)
|
||||
}
|
||||
|
||||
// toKBResponse converts a model to the API response type.
|
||||
func toKBResponse(kb *models.KnowledgeBase) kbResponse {
|
||||
return kbResponse{
|
||||
ID: kb.ID,
|
||||
Name: kb.Name,
|
||||
Description: kb.Description,
|
||||
Scope: kb.Scope,
|
||||
OwnerID: kb.OwnerID,
|
||||
TeamID: kb.TeamID,
|
||||
EmbeddingConfig: kb.EmbeddingConfig,
|
||||
DocumentCount: kb.DocumentCount,
|
||||
ChunkCount: kb.ChunkCount,
|
||||
TotalBytes: kb.TotalBytes,
|
||||
Status: kb.Status,
|
||||
CreatedAt: kb.CreatedAt.Format("2006-01-02T15:04:05Z"),
|
||||
UpdatedAt: kb.UpdatedAt.Format("2006-01-02T15:04:05Z"),
|
||||
}
|
||||
}
|
||||
|
||||
// toDocResponse converts a document model to the API response type.
|
||||
func toDocResponse(doc *models.KBDocument) docResponse {
|
||||
return docResponse{
|
||||
ID: doc.ID,
|
||||
KBID: doc.KBID,
|
||||
Filename: doc.Filename,
|
||||
ContentType: doc.ContentType,
|
||||
SizeBytes: doc.SizeBytes,
|
||||
ChunkCount: doc.ChunkCount,
|
||||
Status: doc.Status,
|
||||
Error: doc.Error,
|
||||
UploadedBy: doc.UploadedBy,
|
||||
CreatedAt: doc.CreatedAt.Format("2006-01-02T15:04:05Z"),
|
||||
UpdatedAt: doc.UpdatedAt.Format("2006-01-02T15:04:05Z"),
|
||||
}
|
||||
}
|
||||
|
||||
// ── Channel KB Toggle ────────────────────────
|
||||
|
||||
// GetChannelKBs returns knowledge bases linked to a channel.
|
||||
// GET /api/v1/channels/:id/knowledge-bases
|
||||
func (h *KnowledgeBaseHandler) GetChannelKBs(c *gin.Context) {
|
||||
channelID := c.Param("id")
|
||||
userID := getUserID(c)
|
||||
|
||||
if !userOwnsChannel(c, channelID, userID) {
|
||||
return
|
||||
}
|
||||
|
||||
kbs, err := h.stores.KnowledgeBases.GetChannelKBs(c.Request.Context(), channelID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to get channel knowledge bases"})
|
||||
return
|
||||
}
|
||||
|
||||
if kbs == nil {
|
||||
kbs = []models.ChannelKB{}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"data": kbs})
|
||||
}
|
||||
|
||||
// SetChannelKBs links knowledge bases to a channel (replaces existing links).
|
||||
// PUT /api/v1/channels/:id/knowledge-bases
|
||||
// Body: { "kb_ids": ["uuid1", "uuid2"] }
|
||||
func (h *KnowledgeBaseHandler) SetChannelKBs(c *gin.Context) {
|
||||
channelID := c.Param("id")
|
||||
userID := getUserID(c)
|
||||
|
||||
if !userOwnsChannel(c, channelID, userID) {
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
KBIDs []string `json:"kb_ids" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Validate that user has access to all requested KBs
|
||||
teamIDs := h.getUserTeamIDs(c, userID)
|
||||
for _, kbID := range req.KBIDs {
|
||||
kb, err := h.stores.KnowledgeBases.GetByID(c.Request.Context(), kbID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("knowledge base %s not found", kbID)})
|
||||
return
|
||||
}
|
||||
if !h.userCanAccess(kb, userID, teamIDs) {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": fmt.Sprintf("no access to knowledge base %s", kbID)})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := h.stores.KnowledgeBases.SetChannelKBs(c.Request.Context(), channelID, req.KBIDs); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to update channel knowledge bases"})
|
||||
return
|
||||
}
|
||||
|
||||
// Return the updated list
|
||||
kbs, _ := h.stores.KnowledgeBases.GetChannelKBs(c.Request.Context(), channelID)
|
||||
if kbs == nil {
|
||||
kbs = []models.ChannelKB{}
|
||||
}
|
||||
|
||||
h.auditLog(c, "kb.channel.link", channelID, map[string]interface{}{
|
||||
"channel_id": channelID, "kb_ids": req.KBIDs,
|
||||
})
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"data": kbs})
|
||||
}
|
||||
|
||||
// userCanAccess checks if a user can access a KB based on scope.
|
||||
func (h *KnowledgeBaseHandler) userCanAccess(kb *models.KnowledgeBase, userID string, teamIDs []string) bool {
|
||||
switch kb.Scope {
|
||||
case "global":
|
||||
return true
|
||||
case "personal":
|
||||
return kb.OwnerID != nil && *kb.OwnerID == userID
|
||||
case "team":
|
||||
if kb.TeamID == nil {
|
||||
return false
|
||||
}
|
||||
for _, tid := range teamIDs {
|
||||
if tid == *kb.TeamID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ── KB Hint for System Prompt ────────────────
|
||||
|
||||
// BuildKBHint returns a system prompt fragment listing active KBs for a
|
||||
// channel, or empty string if none are active. Called by the completion
|
||||
// handler to nudge the LLM to use kb_search.
|
||||
func BuildKBHint(ctx context.Context, stores store.Stores, channelID, userID string) string {
|
||||
teamIDs, _ := stores.Teams.GetUserTeamIDs(ctx, userID)
|
||||
|
||||
kbIDs, err := stores.KnowledgeBases.GetActiveKBIDs(ctx, channelID, userID, teamIDs)
|
||||
if err != nil || len(kbIDs) == 0 {
|
||||
// Also check personal KBs — always available
|
||||
personalKBs, err := stores.KnowledgeBases.ListPersonal(ctx, userID)
|
||||
if err != nil || len(personalKBs) == 0 {
|
||||
return ""
|
||||
}
|
||||
// Only include personal KBs that have chunks
|
||||
hasContent := false
|
||||
for _, kb := range personalKBs {
|
||||
if kb.ChunkCount > 0 {
|
||||
hasContent = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasContent {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// Collect KB details for the hint
|
||||
type kbInfo struct {
|
||||
Name string
|
||||
DocCount int
|
||||
}
|
||||
var kbs []kbInfo
|
||||
|
||||
// Channel-linked KBs
|
||||
channelKBs, err := stores.KnowledgeBases.GetChannelKBs(ctx, channelID)
|
||||
if err == nil {
|
||||
for _, ckb := range channelKBs {
|
||||
if ckb.Enabled && ckb.DocumentCount > 0 {
|
||||
kbs = append(kbs, kbInfo{Name: ckb.KBName, DocCount: ckb.DocumentCount})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Personal KBs (not already linked to channel)
|
||||
personalKBs, err := stores.KnowledgeBases.ListPersonal(ctx, userID)
|
||||
if err == nil {
|
||||
linked := make(map[string]bool)
|
||||
for _, ckb := range channelKBs {
|
||||
linked[ckb.KBID] = true
|
||||
}
|
||||
for _, kb := range personalKBs {
|
||||
if !linked[kb.ID] && kb.ChunkCount > 0 {
|
||||
kbs = append(kbs, kbInfo{Name: kb.Name, DocCount: kb.DocumentCount})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(kbs) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
hint := "\nYou have access to the following knowledge bases:\n"
|
||||
for _, kb := range kbs {
|
||||
hint += fmt.Sprintf("- \"%s\" (%d documents)\n", kb.Name, kb.DocCount)
|
||||
}
|
||||
hint += "Use the kb_search tool to find relevant information from these sources when the user asks questions that might be answered by their documents."
|
||||
|
||||
return hint
|
||||
}
|
||||
|
||||
// ── Audit Logging ──────────────────────────────
|
||||
|
||||
// safeFilename safely extracts the filename from a deleted document.
|
||||
func safeFilename(doc *models.KBDocument) string {
|
||||
if doc == nil {
|
||||
return ""
|
||||
}
|
||||
return doc.Filename
|
||||
}
|
||||
|
||||
// auditLog records a KB operation in the audit log.
|
||||
func (h *KnowledgeBaseHandler) auditLog(c *gin.Context, action, resourceID string, metadata map[string]interface{}) {
|
||||
if h.stores.Audit == nil {
|
||||
return
|
||||
}
|
||||
|
||||
userID := getUserID(c)
|
||||
var meta models.JSONMap
|
||||
if metadata != nil {
|
||||
meta = models.JSONMap(metadata)
|
||||
}
|
||||
|
||||
entry := &models.AuditEntry{
|
||||
ActorID: &userID,
|
||||
Action: action,
|
||||
ResourceType: "knowledge_base",
|
||||
ResourceID: resourceID,
|
||||
Metadata: meta,
|
||||
IPAddress: c.ClientIP(),
|
||||
UserAgent: c.GetHeader("User-Agent"),
|
||||
}
|
||||
|
||||
if err := h.stores.Audit.Log(context.Background(), entry); err != nil {
|
||||
log.Printf("audit log error (KB): %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ── KB Rebuild ─────────────────────────────────
|
||||
|
||||
// RebuildKB re-chunks and re-embeds all documents in a knowledge base.
|
||||
// POST /api/v1/knowledge-bases/:id/rebuild
|
||||
func (h *KnowledgeBaseHandler) RebuildKB(c *gin.Context) {
|
||||
if h.ingester == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "ingestion pipeline not configured"})
|
||||
return
|
||||
}
|
||||
|
||||
kb, ok := h.loadAndAuthorize(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
userID := getUserID(c)
|
||||
|
||||
// List all documents in the KB
|
||||
docs, err := h.stores.KnowledgeBases.ListDocuments(ctx, kb.ID)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list documents"})
|
||||
return
|
||||
}
|
||||
|
||||
if len(docs) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "knowledge base has no documents to rebuild"})
|
||||
return
|
||||
}
|
||||
|
||||
// Resolve team for embedding role
|
||||
teamIDs := h.getUserTeamIDs(c, userID)
|
||||
var teamID *string
|
||||
if len(teamIDs) > 0 {
|
||||
teamID = &teamIDs[0]
|
||||
}
|
||||
|
||||
// Delete all existing chunks (embeddings will be regenerated)
|
||||
rebuiltCount := 0
|
||||
for _, doc := range docs {
|
||||
if doc.ExtractedText == nil || *doc.ExtractedText == "" {
|
||||
// Skip docs without extracted text — they'd need re-extraction
|
||||
// which requires the original file in object store
|
||||
continue
|
||||
}
|
||||
|
||||
// Clear existing chunks for this document
|
||||
if err := h.stores.KnowledgeBases.DeleteChunksForDocument(ctx, doc.ID); err != nil {
|
||||
log.Printf("⚠ rebuild: failed to delete chunks for doc %s: %v", doc.ID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Reset document status and re-ingest
|
||||
h.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "pending", nil)
|
||||
h.ingester.IngestDocument(kb, &doc, userID, teamID)
|
||||
rebuiltCount++
|
||||
}
|
||||
|
||||
h.auditLog(c, "kb.rebuild", kb.ID, map[string]interface{}{
|
||||
"kb_name": kb.Name,
|
||||
"total_docs": len(docs),
|
||||
"rebuilt_docs": rebuiltCount,
|
||||
})
|
||||
|
||||
c.JSON(http.StatusAccepted, gin.H{
|
||||
"message": "rebuild started",
|
||||
"total_docs": len(docs),
|
||||
"rebuilding": rebuiltCount,
|
||||
"skipped": len(docs) - rebuiltCount,
|
||||
})
|
||||
}
|
||||
244
server/knowledge/chunker.go
Normal file
244
server/knowledge/chunker.go
Normal file
@@ -0,0 +1,244 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// ── Configuration ────────────────────────────
|
||||
|
||||
// ChunkConfig controls how text is split into chunks.
|
||||
type ChunkConfig struct {
|
||||
ChunkSize int // target chars per chunk (default 1000)
|
||||
ChunkOverlap int // overlap chars between adjacent chunks (default 200)
|
||||
Separators []string // split hierarchy, tried in order
|
||||
}
|
||||
|
||||
// DefaultChunkConfig returns sensible defaults for general documents.
|
||||
// ~250 tokens per chunk at ~4 chars/token.
|
||||
func DefaultChunkConfig() ChunkConfig {
|
||||
return ChunkConfig{
|
||||
ChunkSize: 1000,
|
||||
ChunkOverlap: 200,
|
||||
Separators: []string{"\n\n", "\n", ". ", " "},
|
||||
}
|
||||
}
|
||||
|
||||
// ── Output ───────────────────────────────────
|
||||
|
||||
// Chunk is a single text segment with position metadata.
|
||||
type Chunk struct {
|
||||
Content string // chunk text
|
||||
Index int // ordinal within document (0-based)
|
||||
TokenCount int // estimated token count (~chars/4)
|
||||
Metadata map[string]any // extensible: page, heading, etc.
|
||||
}
|
||||
|
||||
// ── Splitter ─────────────────────────────────
|
||||
|
||||
// SplitText splits text into overlapping chunks using recursive
|
||||
// character splitting. It tries separators in order, splitting on
|
||||
// the first one that produces segments smaller than ChunkSize.
|
||||
// Segments that are still too large are recursively split with the
|
||||
// next separator.
|
||||
func SplitText(text string, cfg ChunkConfig) []Chunk {
|
||||
if cfg.ChunkSize <= 0 {
|
||||
cfg.ChunkSize = 1000
|
||||
}
|
||||
if cfg.ChunkOverlap < 0 {
|
||||
cfg.ChunkOverlap = 0
|
||||
}
|
||||
if cfg.ChunkOverlap >= cfg.ChunkSize {
|
||||
cfg.ChunkOverlap = cfg.ChunkSize / 5
|
||||
}
|
||||
if len(cfg.Separators) == 0 {
|
||||
cfg.Separators = DefaultChunkConfig().Separators
|
||||
}
|
||||
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
// If text fits in one chunk, return it directly.
|
||||
if utf8.RuneCountInString(text) <= cfg.ChunkSize {
|
||||
return []Chunk{{
|
||||
Content: text,
|
||||
Index: 0,
|
||||
TokenCount: estimateTokens(text),
|
||||
Metadata: map[string]any{},
|
||||
}}
|
||||
}
|
||||
|
||||
// Recursively split, then merge into overlapping windows.
|
||||
segments := recursiveSplit(text, cfg.Separators, cfg.ChunkSize)
|
||||
return mergeSegments(segments, cfg.ChunkSize, cfg.ChunkOverlap)
|
||||
}
|
||||
|
||||
// recursiveSplit tries each separator in order. For the first separator
|
||||
// that actually splits the text, it checks if each piece fits within
|
||||
// chunkSize. Pieces that are still too large recurse with the remaining
|
||||
// separators. Pieces that fit are kept as-is.
|
||||
func recursiveSplit(text string, separators []string, chunkSize int) []string {
|
||||
if utf8.RuneCountInString(text) <= chunkSize {
|
||||
return []string{text}
|
||||
}
|
||||
|
||||
if len(separators) == 0 {
|
||||
// No separators left — hard split by character count.
|
||||
return hardSplit(text, chunkSize)
|
||||
}
|
||||
|
||||
sep := separators[0]
|
||||
rest := separators[1:]
|
||||
|
||||
parts := splitKeepSep(text, sep)
|
||||
if len(parts) <= 1 {
|
||||
// This separator didn't split anything — try next.
|
||||
return recursiveSplit(text, rest, chunkSize)
|
||||
}
|
||||
|
||||
var result []string
|
||||
for _, part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
if utf8.RuneCountInString(part) <= chunkSize {
|
||||
result = append(result, part)
|
||||
} else {
|
||||
// Still too large — recurse with remaining separators.
|
||||
result = append(result, recursiveSplit(part, rest, chunkSize)...)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// splitKeepSep splits text on sep. Unlike strings.Split, the separator
|
||||
// is kept at the end of each segment (except the last) to preserve
|
||||
// context like sentence-ending periods.
|
||||
func splitKeepSep(text, sep string) []string {
|
||||
parts := strings.Split(text, sep)
|
||||
if len(parts) <= 1 {
|
||||
return parts
|
||||
}
|
||||
|
||||
result := make([]string, 0, len(parts))
|
||||
for i, part := range parts {
|
||||
if i < len(parts)-1 {
|
||||
result = append(result, part+sep)
|
||||
} else if part != "" {
|
||||
result = append(result, part)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// hardSplit cuts text into pieces of at most chunkSize runes.
|
||||
// Last resort when no separator works.
|
||||
func hardSplit(text string, chunkSize int) []string {
|
||||
runes := []rune(text)
|
||||
var result []string
|
||||
for i := 0; i < len(runes); i += chunkSize {
|
||||
end := i + chunkSize
|
||||
if end > len(runes) {
|
||||
end = len(runes)
|
||||
}
|
||||
result = append(result, string(runes[i:end]))
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// mergeSegments combines small segments into chunks up to chunkSize,
|
||||
// then applies overlap between adjacent chunks.
|
||||
func mergeSegments(segments []string, chunkSize, overlap int) []Chunk {
|
||||
if len(segments) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// First pass: merge small adjacent segments into windows up to chunkSize.
|
||||
var merged []string
|
||||
var current strings.Builder
|
||||
|
||||
for _, seg := range segments {
|
||||
segLen := utf8.RuneCountInString(seg)
|
||||
curLen := utf8.RuneCountInString(current.String())
|
||||
|
||||
if curLen > 0 && curLen+segLen > chunkSize {
|
||||
// Flush current buffer.
|
||||
merged = append(merged, strings.TrimSpace(current.String()))
|
||||
current.Reset()
|
||||
}
|
||||
if current.Len() > 0 {
|
||||
current.WriteString(" ")
|
||||
}
|
||||
current.WriteString(seg)
|
||||
}
|
||||
if current.Len() > 0 {
|
||||
merged = append(merged, strings.TrimSpace(current.String()))
|
||||
}
|
||||
|
||||
if overlap <= 0 || len(merged) <= 1 {
|
||||
// No overlap needed — convert directly to Chunks.
|
||||
chunks := make([]Chunk, len(merged))
|
||||
for i, text := range merged {
|
||||
chunks[i] = Chunk{
|
||||
Content: text,
|
||||
Index: i,
|
||||
TokenCount: estimateTokens(text),
|
||||
Metadata: map[string]any{},
|
||||
}
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
// Second pass: create overlapping chunks.
|
||||
// Each chunk includes up to `overlap` chars from the end of the
|
||||
// previous chunk as prefix context.
|
||||
chunks := make([]Chunk, 0, len(merged))
|
||||
for i, text := range merged {
|
||||
if i > 0 {
|
||||
prev := merged[i-1]
|
||||
suffix := overlapSuffix(prev, overlap)
|
||||
if suffix != "" {
|
||||
text = suffix + " " + text
|
||||
}
|
||||
}
|
||||
chunks = append(chunks, Chunk{
|
||||
Content: text,
|
||||
Index: i,
|
||||
TokenCount: estimateTokens(text),
|
||||
Metadata: map[string]any{},
|
||||
})
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
// overlapSuffix returns the last `n` characters of s, breaking at a
|
||||
// word boundary to avoid splitting mid-word.
|
||||
func overlapSuffix(s string, n int) string {
|
||||
runes := []rune(s)
|
||||
if len(runes) <= n {
|
||||
return s
|
||||
}
|
||||
|
||||
start := len(runes) - n
|
||||
suffix := string(runes[start:])
|
||||
|
||||
// Try to break at a space to avoid mid-word splits.
|
||||
if idx := strings.Index(suffix, " "); idx >= 0 && idx < len(suffix)/2 {
|
||||
suffix = suffix[idx+1:]
|
||||
}
|
||||
return suffix
|
||||
}
|
||||
|
||||
// estimateTokens gives a rough token count. Most tokenizers average
|
||||
// ~4 characters per token for English text.
|
||||
func estimateTokens(s string) int {
|
||||
n := utf8.RuneCountInString(s)
|
||||
tokens := n / 4
|
||||
if tokens == 0 && n > 0 {
|
||||
tokens = 1
|
||||
}
|
||||
return tokens
|
||||
}
|
||||
227
server/knowledge/chunker_test.go
Normal file
227
server/knowledge/chunker_test.go
Normal file
@@ -0,0 +1,227 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
func TestSplitText_EmptyInput(t *testing.T) {
|
||||
chunks := SplitText("", DefaultChunkConfig())
|
||||
if chunks != nil {
|
||||
t.Errorf("expected nil for empty input, got %d chunks", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_WhitespaceOnly(t *testing.T) {
|
||||
chunks := SplitText(" \n\n ", DefaultChunkConfig())
|
||||
if chunks != nil {
|
||||
t.Errorf("expected nil for whitespace input, got %d chunks", len(chunks))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_SmallText(t *testing.T) {
|
||||
text := "Hello, world."
|
||||
chunks := SplitText(text, DefaultChunkConfig())
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("expected 1 chunk, got %d", len(chunks))
|
||||
}
|
||||
if chunks[0].Content != text {
|
||||
t.Errorf("content mismatch: %q", chunks[0].Content)
|
||||
}
|
||||
if chunks[0].Index != 0 {
|
||||
t.Errorf("expected index 0, got %d", chunks[0].Index)
|
||||
}
|
||||
if chunks[0].TokenCount <= 0 {
|
||||
t.Error("expected positive token count")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_ParagraphSplit(t *testing.T) {
|
||||
// Two paragraphs, each under chunk size, but together over.
|
||||
para := strings.Repeat("word ", 60) // ~300 chars each
|
||||
text := para + "\n\n" + para
|
||||
|
||||
cfg := ChunkConfig{ChunkSize: 400, ChunkOverlap: 0, Separators: []string{"\n\n", "\n", ". ", " "}}
|
||||
chunks := SplitText(text, cfg)
|
||||
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("expected at least 2 chunks from paragraph split, got %d", len(chunks))
|
||||
}
|
||||
|
||||
for i, c := range chunks {
|
||||
if utf8.RuneCountInString(c.Content) == 0 {
|
||||
t.Errorf("chunk %d is empty", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_OverlapPresent(t *testing.T) {
|
||||
// Build text that will split into multiple chunks.
|
||||
sentences := make([]string, 20)
|
||||
for i := range sentences {
|
||||
sentences[i] = "This is sentence number " + strings.Repeat("x", 40) + ". "
|
||||
}
|
||||
text := strings.Join(sentences, "")
|
||||
|
||||
cfg := ChunkConfig{ChunkSize: 200, ChunkOverlap: 50, Separators: []string{". ", " "}}
|
||||
chunks := SplitText(text, cfg)
|
||||
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("expected at least 3 chunks, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Verify chunks have sequential indexes.
|
||||
for i, c := range chunks {
|
||||
if c.Index != i {
|
||||
t.Errorf("chunk %d has index %d", i, c.Index)
|
||||
}
|
||||
}
|
||||
|
||||
// With overlap > 0 and multiple chunks, later chunks should share
|
||||
// some content with the previous chunk's tail.
|
||||
// Just verify the overlap logic ran (chunks 1+ are longer or contain
|
||||
// content from the previous chunk's ending).
|
||||
if len(chunks) >= 2 {
|
||||
// Chunk 1 should contain some overlap from chunk 0's tail.
|
||||
// We can't check exact content easily, but the chunk should be
|
||||
// non-trivially sized.
|
||||
if len(chunks[1].Content) < 50 {
|
||||
t.Errorf("chunk 1 seems too short for overlap: %d chars", len(chunks[1].Content))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_ZeroOverlap(t *testing.T) {
|
||||
text := strings.Repeat("A", 500) + "\n\n" + strings.Repeat("B", 500)
|
||||
cfg := ChunkConfig{ChunkSize: 400, ChunkOverlap: 0, Separators: []string{"\n\n"}}
|
||||
chunks := SplitText(text, cfg)
|
||||
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("expected at least 2 chunks, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Without overlap, chunks should not share content.
|
||||
for i, c := range chunks {
|
||||
if c.Index != i {
|
||||
t.Errorf("chunk %d index = %d", i, c.Index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_HardSplitFallback(t *testing.T) {
|
||||
// No separators at all in the text — forces hard character split.
|
||||
text := strings.Repeat("x", 500)
|
||||
cfg := ChunkConfig{ChunkSize: 200, ChunkOverlap: 0, Separators: []string{"\n\n", "\n", ". ", " "}}
|
||||
chunks := SplitText(text, cfg)
|
||||
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("expected multiple chunks from hard split, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// All chunks should be at most chunkSize.
|
||||
for i, c := range chunks {
|
||||
if utf8.RuneCountInString(c.Content) > cfg.ChunkSize {
|
||||
t.Errorf("chunk %d exceeds chunk size: %d > %d",
|
||||
i, utf8.RuneCountInString(c.Content), cfg.ChunkSize)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_TokenEstimate(t *testing.T) {
|
||||
text := strings.Repeat("word ", 100) // 500 chars
|
||||
chunks := SplitText(text, ChunkConfig{ChunkSize: 2000, ChunkOverlap: 0})
|
||||
|
||||
if len(chunks) != 1 {
|
||||
t.Fatalf("expected 1 chunk, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// ~500 chars / 4 = ~125 tokens
|
||||
if chunks[0].TokenCount < 100 || chunks[0].TokenCount > 150 {
|
||||
t.Errorf("expected ~125 tokens, got %d", chunks[0].TokenCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_InvalidConfig(t *testing.T) {
|
||||
text := "Hello world. This is a test."
|
||||
|
||||
// ChunkSize 0 should use default.
|
||||
chunks := SplitText(text, ChunkConfig{ChunkSize: 0})
|
||||
if len(chunks) != 1 {
|
||||
t.Errorf("expected 1 chunk with zero chunk size (uses default), got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Overlap >= ChunkSize should be clamped.
|
||||
chunks = SplitText(strings.Repeat("word ", 200), ChunkConfig{
|
||||
ChunkSize: 100,
|
||||
ChunkOverlap: 100,
|
||||
})
|
||||
if len(chunks) == 0 {
|
||||
t.Error("expected chunks with overlap == chunk_size (should clamp)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_SentenceSplit(t *testing.T) {
|
||||
text := "First sentence. Second sentence. Third sentence. Fourth sentence. Fifth sentence."
|
||||
cfg := ChunkConfig{ChunkSize: 40, ChunkOverlap: 0, Separators: []string{"\n\n", "\n", ". ", " "}}
|
||||
chunks := SplitText(text, cfg)
|
||||
|
||||
if len(chunks) < 2 {
|
||||
t.Fatalf("expected multiple sentence-split chunks, got %d", len(chunks))
|
||||
}
|
||||
|
||||
// Each chunk should be within bounds.
|
||||
for i, c := range chunks {
|
||||
if utf8.RuneCountInString(c.Content) == 0 {
|
||||
t.Errorf("chunk %d is empty", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitText_MetadataInitialized(t *testing.T) {
|
||||
chunks := SplitText("Some text here.", DefaultChunkConfig())
|
||||
if len(chunks) != 1 {
|
||||
t.Fatal("expected 1 chunk")
|
||||
}
|
||||
if chunks[0].Metadata == nil {
|
||||
t.Error("metadata should be initialized, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateTokens(t *testing.T) {
|
||||
cases := []struct {
|
||||
input string
|
||||
minTok int
|
||||
maxTok int
|
||||
}{
|
||||
{"", 0, 0},
|
||||
{"hi", 1, 1},
|
||||
{"hello world", 2, 4},
|
||||
{strings.Repeat("a", 100), 20, 30},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
got := estimateTokens(tc.input)
|
||||
if got < tc.minTok || got > tc.maxTok {
|
||||
t.Errorf("estimateTokens(%q): got %d, want [%d, %d]", tc.input, got, tc.minTok, tc.maxTok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOverlapSuffix(t *testing.T) {
|
||||
s := "the quick brown fox jumps over the lazy dog"
|
||||
|
||||
// Should return last ~20 chars, breaking at word boundary.
|
||||
suffix := overlapSuffix(s, 20)
|
||||
if len(suffix) == 0 {
|
||||
t.Error("expected non-empty suffix")
|
||||
}
|
||||
if len(suffix) > 25 { // some slack for word boundary adjustment
|
||||
t.Errorf("suffix too long: %d chars: %q", len(suffix), suffix)
|
||||
}
|
||||
|
||||
// Short string: return whole thing.
|
||||
short := "hello"
|
||||
if got := overlapSuffix(short, 100); got != short {
|
||||
t.Errorf("expected %q, got %q", short, got)
|
||||
}
|
||||
}
|
||||
184
server/knowledge/embedder.go
Normal file
184
server/knowledge/embedder.go
Normal file
@@ -0,0 +1,184 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/models"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/roles"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/store"
|
||||
)
|
||||
|
||||
// ── Configuration ────────────────────────────
|
||||
|
||||
const (
|
||||
// DefaultBatchSize is the max chunks per embedding API call.
|
||||
// Most providers handle 100 inputs per request comfortably.
|
||||
DefaultBatchSize = 100
|
||||
|
||||
// MaxVectorDim is the column width in pgvector (vector(3072)).
|
||||
// Vectors shorter than this are zero-padded on storage.
|
||||
MaxVectorDim = 3072
|
||||
)
|
||||
|
||||
// ── Embedder ─────────────────────────────────
|
||||
|
||||
// Embedder generates vector embeddings for text chunks using the
|
||||
// embedding role resolver. It handles batching and dimension
|
||||
// normalization (zero-padding to MaxVectorDim).
|
||||
type Embedder struct {
|
||||
resolver *roles.Resolver
|
||||
stores store.Stores // optional — for usage tracking
|
||||
batchSize int
|
||||
}
|
||||
|
||||
// NewEmbedder creates an embedder that dispatches through the role resolver.
|
||||
func NewEmbedder(resolver *roles.Resolver) *Embedder {
|
||||
return &Embedder{
|
||||
resolver: resolver,
|
||||
batchSize: DefaultBatchSize,
|
||||
}
|
||||
}
|
||||
|
||||
// WithStores attaches stores for usage tracking. Returns self for chaining.
|
||||
func (e *Embedder) WithStores(s store.Stores) *Embedder {
|
||||
e.stores = s
|
||||
return e
|
||||
}
|
||||
|
||||
// EmbedResult holds the output of a batch embedding operation.
|
||||
type EmbedResult struct {
|
||||
Vectors [][]float64 // one per input chunk, zero-padded to MaxVectorDim
|
||||
Model string // model that produced the embeddings
|
||||
Dimensions int // native dimension before padding
|
||||
InputTokens int // total tokens consumed across all batches
|
||||
ConfigID string // provider config that was used
|
||||
ProviderScope string // "personal", "team", or "global"
|
||||
}
|
||||
|
||||
// EmbedChunks generates embeddings for a slice of text chunks.
|
||||
// Chunks are batched in groups of batchSize to avoid provider limits.
|
||||
// Uses the embedding role resolution chain: personal → team → global.
|
||||
//
|
||||
// Returns vectors in the same order as input chunks.
|
||||
func (e *Embedder) EmbedChunks(ctx context.Context, userID string, teamID *string, texts []string) (*EmbedResult, error) {
|
||||
if len(texts) == 0 {
|
||||
return &EmbedResult{Vectors: [][]float64{}}, nil
|
||||
}
|
||||
|
||||
// Verify embedding role is configured before starting
|
||||
cfg, err := e.resolver.GetConfig(ctx, roles.RoleEmbedding, userID, teamID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embedding role not configured: %w", err)
|
||||
}
|
||||
if cfg.Primary == nil {
|
||||
return nil, fmt.Errorf("embedding role has no primary binding")
|
||||
}
|
||||
|
||||
allVectors := make([][]float64, 0, len(texts))
|
||||
var model string
|
||||
var nativeDim int
|
||||
var configID, provScope string
|
||||
totalTokens := 0
|
||||
|
||||
for i := 0; i < len(texts); i += e.batchSize {
|
||||
end := i + e.batchSize
|
||||
if end > len(texts) {
|
||||
end = len(texts)
|
||||
}
|
||||
batch := texts[i:end]
|
||||
|
||||
log.Printf(" embed batch %d/%d (%d chunks)",
|
||||
(i/e.batchSize)+1, (len(texts)+e.batchSize-1)/e.batchSize, len(batch))
|
||||
|
||||
result, err := e.resolver.Embed(ctx, roles.RoleEmbedding, userID, teamID, batch)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("embed batch starting at chunk %d: %w", i, err)
|
||||
}
|
||||
|
||||
if len(result.Embeddings) != len(batch) {
|
||||
return nil, fmt.Errorf("embed batch %d: expected %d vectors, got %d",
|
||||
i/e.batchSize, len(batch), len(result.Embeddings))
|
||||
}
|
||||
|
||||
// Capture model info from first batch
|
||||
if model == "" {
|
||||
model = result.Model
|
||||
configID = result.ConfigID
|
||||
provScope = result.ProviderScope
|
||||
if len(result.Embeddings) > 0 {
|
||||
nativeDim = len(result.Embeddings[0])
|
||||
}
|
||||
}
|
||||
|
||||
totalTokens += result.InputTokens
|
||||
allVectors = append(allVectors, result.Embeddings...)
|
||||
}
|
||||
|
||||
// Zero-pad to MaxVectorDim for uniform storage
|
||||
for i, vec := range allVectors {
|
||||
allVectors[i] = padVector(vec, MaxVectorDim)
|
||||
}
|
||||
|
||||
return &EmbedResult{
|
||||
Vectors: allVectors,
|
||||
Model: model,
|
||||
Dimensions: nativeDim,
|
||||
InputTokens: totalTokens,
|
||||
ConfigID: configID,
|
||||
ProviderScope: provScope,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LogUsage records an embedding usage entry. Silently no-ops if stores
|
||||
// were not attached via WithStores or if there are no tokens to log.
|
||||
func (e *Embedder) LogUsage(ctx context.Context, userID string, channelID *string, result *EmbedResult) {
|
||||
if e.stores.Usage == nil || result == nil || result.InputTokens == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
role := "embedding"
|
||||
entry := &models.UsageEntry{
|
||||
ChannelID: channelID,
|
||||
UserID: userID,
|
||||
ProviderConfigID: strPtr(result.ConfigID),
|
||||
ProviderScope: result.ProviderScope,
|
||||
ModelID: result.Model,
|
||||
Role: &role,
|
||||
PromptTokens: result.InputTokens,
|
||||
CompletionTokens: 0,
|
||||
}
|
||||
|
||||
if err := e.stores.Usage.Log(ctx, entry); err != nil {
|
||||
log.Printf("⚠ Failed to log embedding usage: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// IsConfigured returns true if the embedding role has at least a primary binding.
|
||||
func (e *Embedder) IsConfigured(ctx context.Context) bool {
|
||||
return e.resolver.IsConfigured(ctx, roles.RoleEmbedding)
|
||||
}
|
||||
|
||||
// ── Helpers ──────────────────────────────────
|
||||
|
||||
// padVector zero-pads a vector to the target dimension.
|
||||
// If the vector is already the target size or larger, it's truncated.
|
||||
func padVector(vec []float64, targetDim int) []float64 {
|
||||
if len(vec) == targetDim {
|
||||
return vec
|
||||
}
|
||||
if len(vec) > targetDim {
|
||||
return vec[:targetDim]
|
||||
}
|
||||
padded := make([]float64, targetDim)
|
||||
copy(padded, vec)
|
||||
return padded
|
||||
}
|
||||
|
||||
func strPtr(s string) *string {
|
||||
if s == "" {
|
||||
return nil
|
||||
}
|
||||
return &s
|
||||
}
|
||||
258
server/knowledge/ingest.go
Normal file
258
server/knowledge/ingest.go
Normal file
@@ -0,0 +1,258 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/models"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/storage"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/store"
|
||||
)
|
||||
|
||||
// ── Configuration ────────────────────────────
|
||||
|
||||
const (
|
||||
// DefaultConcurrency limits parallel ingestion goroutines.
|
||||
DefaultConcurrency = 3
|
||||
|
||||
// MaxTextLength is the upper bound for extracted text (10 MB).
|
||||
// Documents larger than this are truncated with a warning.
|
||||
MaxTextLength = 10 * 1024 * 1024
|
||||
|
||||
// contextTimeout for the entire ingestion of one document.
|
||||
ingestTimeout = 10 * time.Minute
|
||||
)
|
||||
|
||||
// textMIMETypes are content types we can extract text from directly.
|
||||
var textMIMETypes = map[string]bool{
|
||||
"text/plain": true,
|
||||
"text/markdown": true,
|
||||
"text/csv": true,
|
||||
"text/html": true,
|
||||
}
|
||||
|
||||
// ── Ingester ─────────────────────────────────
|
||||
|
||||
// Ingester manages the document ingestion pipeline.
|
||||
// It runs chunk + embed as background goroutines with a concurrency
|
||||
// semaphore to avoid overwhelming the embedding provider.
|
||||
type Ingester struct {
|
||||
stores store.Stores
|
||||
embedder *Embedder
|
||||
objStore storage.ObjectStore
|
||||
sem chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewIngester creates an ingestion pipeline.
|
||||
func NewIngester(stores store.Stores, embedder *Embedder, objStore storage.ObjectStore, concurrency int) *Ingester {
|
||||
if concurrency <= 0 {
|
||||
concurrency = DefaultConcurrency
|
||||
}
|
||||
return &Ingester{
|
||||
stores: stores,
|
||||
embedder: embedder,
|
||||
objStore: objStore,
|
||||
sem: make(chan struct{}, concurrency),
|
||||
}
|
||||
}
|
||||
|
||||
// IngestDocument starts async ingestion for a document.
|
||||
// The document must already exist in kb_documents with status "pending"
|
||||
// and the file must be stored at doc.StorageKey in the object store.
|
||||
//
|
||||
// The goroutine progresses through: pending → extracting → chunking →
|
||||
// embedding → ready (or → error on failure).
|
||||
func (ing *Ingester) IngestDocument(kb *models.KnowledgeBase, doc *models.KBDocument, userID string, teamID *string) {
|
||||
ing.wg.Add(1)
|
||||
go func() {
|
||||
defer ing.wg.Done()
|
||||
|
||||
// Acquire semaphore slot
|
||||
ing.sem <- struct{}{}
|
||||
defer func() { <-ing.sem }()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), ingestTimeout)
|
||||
defer cancel()
|
||||
|
||||
if err := ing.ingest(ctx, kb, doc, userID, teamID); err != nil {
|
||||
errMsg := err.Error()
|
||||
log.Printf("❌ ingest doc %s (%s): %v", doc.ID, doc.Filename, err)
|
||||
ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "error", &errMsg)
|
||||
ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{"status": "active"})
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Wait blocks until all in-flight ingestion goroutines complete.
|
||||
// Call during graceful shutdown.
|
||||
func (ing *Ingester) Wait() {
|
||||
ing.wg.Wait()
|
||||
}
|
||||
|
||||
// ── Internal Pipeline ────────────────────────
|
||||
|
||||
func (ing *Ingester) ingest(ctx context.Context, kb *models.KnowledgeBase, doc *models.KBDocument, userID string, teamID *string) error {
|
||||
start := time.Now()
|
||||
log.Printf("▶ ingest %s (%s, %d bytes)", doc.Filename, doc.ContentType, doc.SizeBytes)
|
||||
|
||||
// Mark KB as processing
|
||||
ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{"status": "processing"})
|
||||
|
||||
// ── Step 1: Extract text ────────────────
|
||||
ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "extracting", nil)
|
||||
|
||||
text, err := ing.extractText(ctx, doc)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract: %w", err)
|
||||
}
|
||||
|
||||
if strings.TrimSpace(text) == "" {
|
||||
return fmt.Errorf("extract: no text content found in %s", doc.Filename)
|
||||
}
|
||||
|
||||
// Truncate very large documents
|
||||
if len(text) > MaxTextLength {
|
||||
log.Printf(" ⚠ truncating %s from %d to %d chars", doc.Filename, len(text), MaxTextLength)
|
||||
text = text[:MaxTextLength]
|
||||
}
|
||||
|
||||
// Store extracted text for potential re-chunking later
|
||||
if err := ing.stores.KnowledgeBases.UpdateDocumentText(ctx, doc.ID, text, 0); err != nil {
|
||||
log.Printf(" ⚠ failed to store extracted text for %s: %v", doc.ID, err)
|
||||
}
|
||||
|
||||
// ── Step 2: Chunk ───────────────────────
|
||||
ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "chunking", nil)
|
||||
|
||||
cfg := DefaultChunkConfig()
|
||||
chunks := SplitText(text, cfg)
|
||||
if len(chunks) == 0 {
|
||||
return fmt.Errorf("chunk: produced zero chunks from %s", doc.Filename)
|
||||
}
|
||||
|
||||
log.Printf(" ✓ chunked %s → %d chunks", doc.Filename, len(chunks))
|
||||
|
||||
// ── Step 3: Embed ───────────────────────
|
||||
ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "embedding", nil)
|
||||
|
||||
texts := make([]string, len(chunks))
|
||||
for i, c := range chunks {
|
||||
texts[i] = c.Content
|
||||
}
|
||||
|
||||
embedResult, err := ing.embedder.EmbedChunks(ctx, userID, teamID, texts)
|
||||
if err != nil {
|
||||
return fmt.Errorf("embed: %w", err)
|
||||
}
|
||||
|
||||
log.Printf(" ✓ embedded %d chunks (model=%s, dim=%d, tokens=%d)",
|
||||
len(chunks), embedResult.Model, embedResult.Dimensions, embedResult.InputTokens)
|
||||
|
||||
// Track embedding usage (role = "embedding")
|
||||
ing.embedder.LogUsage(ctx, userID, nil, embedResult)
|
||||
|
||||
// ── Step 4: Store chunks with vectors ───
|
||||
dbChunks := make([]models.KBChunk, len(chunks))
|
||||
for i, c := range chunks {
|
||||
dbChunks[i] = models.KBChunk{
|
||||
KBID: kb.ID,
|
||||
DocumentID: doc.ID,
|
||||
ChunkIndex: c.Index,
|
||||
Content: c.Content,
|
||||
TokenCount: c.TokenCount,
|
||||
Embedding: embedResult.Vectors[i],
|
||||
Metadata: models.JSONMap{
|
||||
"char_start": i * (cfg.ChunkSize - cfg.ChunkOverlap), // approximate
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
if err := ing.stores.KnowledgeBases.InsertChunks(ctx, dbChunks); err != nil {
|
||||
return fmt.Errorf("store chunks: %w", err)
|
||||
}
|
||||
|
||||
// ── Step 5: Update stats ────────────────
|
||||
// Update document status + chunk count
|
||||
if err := ing.stores.KnowledgeBases.UpdateDocumentText(ctx, doc.ID, text, len(chunks)); err != nil {
|
||||
log.Printf(" ⚠ failed to update chunk count for %s: %v", doc.ID, err)
|
||||
}
|
||||
ing.stores.KnowledgeBases.UpdateDocumentStatus(ctx, doc.ID, "ready", nil)
|
||||
|
||||
// Update KB-level stats
|
||||
if err := ing.stores.KnowledgeBases.UpdateStats(ctx, kb.ID); err != nil {
|
||||
log.Printf(" ⚠ failed to update KB stats: %v", err)
|
||||
}
|
||||
|
||||
// Store embedding model info in KB config if not set
|
||||
if kb.EmbeddingConfig == nil || kb.EmbeddingConfig["model"] == nil {
|
||||
ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{
|
||||
"embedding_config": models.JSONMap{
|
||||
"model": embedResult.Model,
|
||||
"dimensions": embedResult.Dimensions,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// Mark KB as active
|
||||
ing.stores.KnowledgeBases.Update(ctx, kb.ID, map[string]interface{}{"status": "active"})
|
||||
|
||||
log.Printf(" ✓ ingest complete: %s (%d chunks, %s)",
|
||||
doc.Filename, len(chunks), time.Since(start).Round(time.Millisecond))
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ── Text Extraction ──────────────────────────
|
||||
|
||||
// extractText reads the document file and returns its text content.
|
||||
// For text-based files, reads directly from object store.
|
||||
// For complex formats (PDF, docx), returns an error — the extraction
|
||||
// sidecar integration is planned for a future phase.
|
||||
func (ing *Ingester) extractText(ctx context.Context, doc *models.KBDocument) (string, error) {
|
||||
// If extracted text was already stored (e.g. from a rebuild), use it
|
||||
if doc.ExtractedText != nil && *doc.ExtractedText != "" {
|
||||
return *doc.ExtractedText, nil
|
||||
}
|
||||
|
||||
// Text-based files: read directly
|
||||
if textMIMETypes[doc.ContentType] {
|
||||
return ing.readTextFile(ctx, doc.StorageKey)
|
||||
}
|
||||
|
||||
// Complex document types are not yet supported for inline extraction.
|
||||
// The extraction sidecar (v0.12.0) handles attachments but isn't yet
|
||||
// wired into the KB pipeline. Planned for a future phase.
|
||||
return "", fmt.Errorf("unsupported content type %q — text-based files only for now (txt, md, csv, html)", doc.ContentType)
|
||||
}
|
||||
|
||||
// readTextFile reads a text file from the object store and returns its contents.
|
||||
func (ing *Ingester) readTextFile(ctx context.Context, storageKey string) (string, error) {
|
||||
rc, size, _, err := ing.objStore.Get(ctx, storageKey)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read %s: %w", storageKey, err)
|
||||
}
|
||||
defer rc.Close()
|
||||
|
||||
// Guard against absurdly large files
|
||||
if size > MaxTextLength+1024 {
|
||||
// Read up to limit
|
||||
limited := io.LimitReader(rc, MaxTextLength)
|
||||
data, err := io.ReadAll(limited)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read %s: %w", storageKey, err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
data, err := io.ReadAll(rc)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read %s: %w", storageKey, err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
@@ -17,13 +17,14 @@ import (
|
||||
"git.gobha.me/xcaliber/chat-switchboard/events"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/extraction"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/handlers"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/knowledge"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/middleware"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/providers"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/roles"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/storage"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/store"
|
||||
postgres "git.gobha.me/xcaliber/chat-switchboard/store/postgres"
|
||||
_ "git.gobha.me/xcaliber/chat-switchboard/tools" // registers built-in tools via init()
|
||||
"git.gobha.me/xcaliber/chat-switchboard/tools"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/tools/search"
|
||||
)
|
||||
|
||||
@@ -146,6 +147,19 @@ func main() {
|
||||
// Role resolver for model role dispatch (needs stores + vault)
|
||||
roleResolver := roles.NewResolver(stores, keyResolver)
|
||||
|
||||
// ── Knowledge Base Pipeline ─────────────
|
||||
// Embedder + ingester for RAG document processing (v0.14.0).
|
||||
// Nil-safe: handler checks embedder.IsConfigured() before accepting uploads.
|
||||
kbEmbedder := knowledge.NewEmbedder(roleResolver).WithStores(stores)
|
||||
kbIngester := knowledge.NewIngester(stores, kbEmbedder, objStore, knowledge.DefaultConcurrency)
|
||||
defer kbIngester.Wait() // drain in-flight ingestions on shutdown
|
||||
|
||||
// Register kb_search tool (late registration — needs stores + embedder)
|
||||
tools.RegisterKBSearch(stores, kbEmbedder)
|
||||
|
||||
// Register note tools (late registration — needs stores + embedder for semantic search)
|
||||
tools.RegisterNoteTools(stores, kbEmbedder)
|
||||
|
||||
r := gin.Default()
|
||||
r.Use(middleware.CORS())
|
||||
|
||||
@@ -304,6 +318,22 @@ func main() {
|
||||
// Hook: clean up storage files when channels are deleted
|
||||
handlers.SetChannelDeleteHook(attachH.CleanupChannelStorage)
|
||||
|
||||
// Knowledge Bases (RAG — v0.14.0)
|
||||
kbH := handlers.NewKnowledgeBaseHandler(stores, objStore, kbIngester, kbEmbedder)
|
||||
protected.POST("/knowledge-bases", kbH.CreateKB)
|
||||
protected.GET("/knowledge-bases", kbH.ListKBs)
|
||||
protected.GET("/knowledge-bases/:id", kbH.GetKB)
|
||||
protected.PUT("/knowledge-bases/:id", kbH.UpdateKB)
|
||||
protected.DELETE("/knowledge-bases/:id", kbH.DeleteKB)
|
||||
protected.POST("/knowledge-bases/:id/documents", kbH.UploadDocument)
|
||||
protected.GET("/knowledge-bases/:id/documents", kbH.ListDocuments)
|
||||
protected.GET("/knowledge-bases/:id/documents/:docId/status", kbH.GetDocumentStatus)
|
||||
protected.DELETE("/knowledge-bases/:id/documents/:docId", kbH.DeleteDocument)
|
||||
protected.POST("/knowledge-bases/:id/search", kbH.SearchKB)
|
||||
protected.POST("/knowledge-bases/:id/rebuild", kbH.RebuildKB)
|
||||
protected.GET("/channels/:id/knowledge-bases", kbH.GetChannelKBs)
|
||||
protected.PUT("/channels/:id/knowledge-bases", kbH.SetChannelKBs)
|
||||
|
||||
// Teams (user: my teams)
|
||||
teams := handlers.NewTeamHandler(keyResolver)
|
||||
protected.GET("/teams/mine", teams.MyTeams)
|
||||
|
||||
@@ -622,3 +622,69 @@ func StringPtr(s string) *string { return &s }
|
||||
func Float64Ptr(f float64) *float64 { return &f }
|
||||
func IntPtr(i int) *int { return &i }
|
||||
func BoolPtr(b bool) *bool { return &b }
|
||||
|
||||
// ── Knowledge Bases ────────────────────────────
|
||||
|
||||
// KnowledgeBase is a named collection of documents with vector embeddings.
|
||||
type KnowledgeBase struct {
|
||||
ID string `json:"id" db:"id"`
|
||||
Name string `json:"name" db:"name"`
|
||||
Description string `json:"description" db:"description"`
|
||||
Scope string `json:"scope" db:"scope"` // global, team, personal
|
||||
OwnerID *string `json:"owner_id,omitempty" db:"owner_id"`
|
||||
TeamID *string `json:"team_id,omitempty" db:"team_id"`
|
||||
EmbeddingConfig JSONMap `json:"embedding_config" db:"embedding_config"`
|
||||
DocumentCount int `json:"document_count" db:"document_count"`
|
||||
ChunkCount int `json:"chunk_count" db:"chunk_count"`
|
||||
TotalBytes int64 `json:"total_bytes" db:"total_bytes"`
|
||||
Status string `json:"status" db:"status"` // active, processing, error
|
||||
CreatedAt time.Time `json:"created_at" db:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at" db:"updated_at"`
|
||||
}
|
||||
|
||||
// KBDocument is a single uploaded file within a knowledge base.
|
||||
type KBDocument struct {
|
||||
ID string `json:"id" db:"id"`
|
||||
KBID string `json:"kb_id" db:"kb_id"`
|
||||
Filename string `json:"filename" db:"filename"`
|
||||
ContentType string `json:"content_type" db:"content_type"`
|
||||
SizeBytes int64 `json:"size_bytes" db:"size_bytes"`
|
||||
StorageKey string `json:"storage_key" db:"storage_key"`
|
||||
ExtractedText *string `json:"-" db:"extracted_text"`
|
||||
ChunkCount int `json:"chunk_count" db:"chunk_count"`
|
||||
Status string `json:"status" db:"status"` // pending, chunking, embedding, ready, error
|
||||
Error *string `json:"error,omitempty" db:"error"`
|
||||
UploadedBy string `json:"uploaded_by" db:"uploaded_by"`
|
||||
CreatedAt time.Time `json:"created_at" db:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at" db:"updated_at"`
|
||||
}
|
||||
|
||||
// KBChunk is a text segment from a document with its embedding vector.
|
||||
type KBChunk struct {
|
||||
ID string `json:"id" db:"id"`
|
||||
KBID string `json:"kb_id" db:"kb_id"`
|
||||
DocumentID string `json:"document_id" db:"document_id"`
|
||||
ChunkIndex int `json:"chunk_index" db:"chunk_index"`
|
||||
Content string `json:"content" db:"content"`
|
||||
TokenCount int `json:"token_count" db:"token_count"`
|
||||
Embedding []float64 `json:"-" db:"embedding"` // never serialized to API
|
||||
Metadata JSONMap `json:"metadata" db:"metadata"`
|
||||
CreatedAt time.Time `json:"created_at" db:"created_at"`
|
||||
}
|
||||
|
||||
// KBSearchResult is a single result from similarity search.
|
||||
type KBSearchResult struct {
|
||||
Content string `json:"content"`
|
||||
Filename string `json:"source"`
|
||||
KBName string `json:"kb"`
|
||||
Similarity float64 `json:"similarity"`
|
||||
Metadata JSONMap `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
// ChannelKB represents a knowledge base linked to a channel.
|
||||
type ChannelKB struct {
|
||||
KBID string `json:"kb_id" db:"kb_id"`
|
||||
KBName string `json:"kb_name" db:"name"`
|
||||
Enabled bool `json:"enabled" db:"enabled"`
|
||||
DocumentCount int `json:"document_count" db:"document_count"`
|
||||
}
|
||||
|
||||
@@ -34,6 +34,7 @@ type Stores struct {
|
||||
Pricing PricingStore
|
||||
Extensions ExtensionStore
|
||||
Attachments AttachmentStore
|
||||
KnowledgeBases KnowledgeBaseStore
|
||||
}
|
||||
|
||||
// =========================================
|
||||
@@ -359,6 +360,47 @@ type AttachmentStore interface {
|
||||
ListOrphans(ctx context.Context, olderThan time.Duration) ([]models.Attachment, error)
|
||||
}
|
||||
|
||||
// =========================================
|
||||
// KNOWLEDGE BASE STORE
|
||||
// =========================================
|
||||
|
||||
type KnowledgeBaseStore interface {
|
||||
// KB CRUD
|
||||
Create(ctx context.Context, kb *models.KnowledgeBase) error
|
||||
GetByID(ctx context.Context, id string) (*models.KnowledgeBase, error)
|
||||
Update(ctx context.Context, id string, fields map[string]interface{}) error
|
||||
Delete(ctx context.Context, id string) error
|
||||
|
||||
// Scoped listing
|
||||
ListForUser(ctx context.Context, userID string, teamIDs []string) ([]models.KnowledgeBase, error)
|
||||
ListGlobal(ctx context.Context) ([]models.KnowledgeBase, error)
|
||||
ListForTeam(ctx context.Context, teamID string) ([]models.KnowledgeBase, error)
|
||||
ListPersonal(ctx context.Context, userID string) ([]models.KnowledgeBase, error)
|
||||
|
||||
// Documents
|
||||
CreateDocument(ctx context.Context, doc *models.KBDocument) error
|
||||
GetDocument(ctx context.Context, id string) (*models.KBDocument, error)
|
||||
ListDocuments(ctx context.Context, kbID string) ([]models.KBDocument, error)
|
||||
UpdateDocumentStatus(ctx context.Context, id string, status string, errMsg *string) error
|
||||
UpdateDocumentText(ctx context.Context, id string, text string, chunkCount int) error
|
||||
DeleteDocument(ctx context.Context, id string) (*models.KBDocument, error) // returns for storage cleanup
|
||||
|
||||
// Chunks
|
||||
InsertChunks(ctx context.Context, chunks []models.KBChunk) error
|
||||
DeleteChunksForDocument(ctx context.Context, documentID string) error
|
||||
SimilaritySearch(ctx context.Context, kbIDs []string, queryVec []float64,
|
||||
threshold float64, limit int) ([]models.KBSearchResult, error)
|
||||
|
||||
// Channel links
|
||||
SetChannelKBs(ctx context.Context, channelID string, kbIDs []string) error
|
||||
GetChannelKBs(ctx context.Context, channelID string) ([]models.ChannelKB, error)
|
||||
GetActiveKBIDs(ctx context.Context, channelID string, userID string,
|
||||
teamIDs []string) ([]string, error) // enabled + user has access
|
||||
|
||||
// Stats — recount document_count/chunk_count/total_bytes from child tables
|
||||
UpdateStats(ctx context.Context, kbID string) error
|
||||
}
|
||||
|
||||
// =========================================
|
||||
// SHARED TYPES
|
||||
// =========================================
|
||||
|
||||
432
server/store/postgres/knowledge_bases.go
Normal file
432
server/store/postgres/knowledge_bases.go
Normal file
@@ -0,0 +1,432 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/models"
|
||||
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
// ── KnowledgeBaseStore ──────────────────────────
|
||||
|
||||
type KnowledgeBaseStore struct{}
|
||||
|
||||
func NewKnowledgeBaseStore() *KnowledgeBaseStore { return &KnowledgeBaseStore{} }
|
||||
|
||||
// ── KB CRUD ──────────────────────────────────────
|
||||
|
||||
func (s *KnowledgeBaseStore) Create(ctx context.Context, kb *models.KnowledgeBase) error {
|
||||
return DB.QueryRowContext(ctx, `
|
||||
INSERT INTO knowledge_bases (name, description, scope, owner_id, team_id, embedding_config, status)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id, document_count, chunk_count, total_bytes, created_at, updated_at`,
|
||||
kb.Name, kb.Description, kb.Scope,
|
||||
models.NullString(kb.OwnerID), models.NullString(kb.TeamID),
|
||||
ToJSON(kb.EmbeddingConfig), kb.Status,
|
||||
).Scan(&kb.ID, &kb.DocumentCount, &kb.ChunkCount, &kb.TotalBytes, &kb.CreatedAt, &kb.UpdatedAt)
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) GetByID(ctx context.Context, id string) (*models.KnowledgeBase, error) {
|
||||
kb, err := scanKB(DB.QueryRowContext(ctx, `
|
||||
SELECT id, name, description, scope, owner_id, team_id, embedding_config,
|
||||
document_count, chunk_count, total_bytes, status, created_at, updated_at
|
||||
FROM knowledge_bases WHERE id = $1`, id))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return kb, nil
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) Update(ctx context.Context, id string, fields map[string]interface{}) error {
|
||||
b := NewUpdate("knowledge_bases")
|
||||
for k, v := range fields {
|
||||
if k == "embedding_config" {
|
||||
b.SetJSON(k, v)
|
||||
} else {
|
||||
b.Set(k, v)
|
||||
}
|
||||
}
|
||||
b.Set("updated_at", "now()")
|
||||
if !b.HasSets() {
|
||||
return nil
|
||||
}
|
||||
b.Where("id", id)
|
||||
_, err := b.Exec(DB)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) Delete(ctx context.Context, id string) error {
|
||||
// CASCADE deletes kb_documents, kb_chunks, channel_knowledge_bases.
|
||||
_, err := DB.ExecContext(ctx, "DELETE FROM knowledge_bases WHERE id = $1", id)
|
||||
return err
|
||||
}
|
||||
|
||||
// ── Scoped Listing ──────────────────────────────
|
||||
|
||||
func (s *KnowledgeBaseStore) ListForUser(ctx context.Context, userID string, teamIDs []string) ([]models.KnowledgeBase, error) {
|
||||
// User can see: global + their teams + personal.
|
||||
q := `
|
||||
SELECT id, name, description, scope, owner_id, team_id, embedding_config,
|
||||
document_count, chunk_count, total_bytes, status, created_at, updated_at
|
||||
FROM knowledge_bases
|
||||
WHERE scope = 'global'
|
||||
OR (scope = 'personal' AND owner_id = $1)`
|
||||
|
||||
args := []interface{}{userID}
|
||||
|
||||
if len(teamIDs) > 0 {
|
||||
q += fmt.Sprintf(` OR (scope = 'team' AND team_id = ANY($%d))`, len(args)+1)
|
||||
args = append(args, pq.Array(teamIDs))
|
||||
}
|
||||
|
||||
q += ` ORDER BY name`
|
||||
return queryKBs(ctx, q, args...)
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) ListGlobal(ctx context.Context) ([]models.KnowledgeBase, error) {
|
||||
return queryKBs(ctx, `
|
||||
SELECT id, name, description, scope, owner_id, team_id, embedding_config,
|
||||
document_count, chunk_count, total_bytes, status, created_at, updated_at
|
||||
FROM knowledge_bases WHERE scope = 'global' ORDER BY name`)
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) ListForTeam(ctx context.Context, teamID string) ([]models.KnowledgeBase, error) {
|
||||
return queryKBs(ctx, `
|
||||
SELECT id, name, description, scope, owner_id, team_id, embedding_config,
|
||||
document_count, chunk_count, total_bytes, status, created_at, updated_at
|
||||
FROM knowledge_bases WHERE team_id = $1 ORDER BY name`, teamID)
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) ListPersonal(ctx context.Context, userID string) ([]models.KnowledgeBase, error) {
|
||||
return queryKBs(ctx, `
|
||||
SELECT id, name, description, scope, owner_id, team_id, embedding_config,
|
||||
document_count, chunk_count, total_bytes, status, created_at, updated_at
|
||||
FROM knowledge_bases WHERE scope = 'personal' AND owner_id = $1 ORDER BY name`, userID)
|
||||
}
|
||||
|
||||
// ── Documents ────────────────────────────────────
|
||||
|
||||
func (s *KnowledgeBaseStore) CreateDocument(ctx context.Context, doc *models.KBDocument) error {
|
||||
return DB.QueryRowContext(ctx, `
|
||||
INSERT INTO kb_documents (kb_id, filename, content_type, size_bytes, storage_key, status, uploaded_by)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7)
|
||||
RETURNING id, chunk_count, created_at, updated_at`,
|
||||
doc.KBID, doc.Filename, doc.ContentType, doc.SizeBytes,
|
||||
doc.StorageKey, doc.Status, doc.UploadedBy,
|
||||
).Scan(&doc.ID, &doc.ChunkCount, &doc.CreatedAt, &doc.UpdatedAt)
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) GetDocument(ctx context.Context, id string) (*models.KBDocument, error) {
|
||||
var doc models.KBDocument
|
||||
var extractedText, errMsg sql.NullString
|
||||
err := DB.QueryRowContext(ctx, `
|
||||
SELECT id, kb_id, filename, content_type, size_bytes, storage_key,
|
||||
extracted_text, chunk_count, status, error, uploaded_by, created_at, updated_at
|
||||
FROM kb_documents WHERE id = $1`, id).Scan(
|
||||
&doc.ID, &doc.KBID, &doc.Filename, &doc.ContentType, &doc.SizeBytes,
|
||||
&doc.StorageKey, &extractedText, &doc.ChunkCount, &doc.Status,
|
||||
&errMsg, &doc.UploadedBy, &doc.CreatedAt, &doc.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
doc.ExtractedText = NullableStringPtr(extractedText)
|
||||
doc.Error = NullableStringPtr(errMsg)
|
||||
return &doc, nil
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) ListDocuments(ctx context.Context, kbID string) ([]models.KBDocument, error) {
|
||||
rows, err := DB.QueryContext(ctx, `
|
||||
SELECT id, kb_id, filename, content_type, size_bytes, storage_key,
|
||||
chunk_count, status, error, uploaded_by, created_at, updated_at
|
||||
FROM kb_documents WHERE kb_id = $1 ORDER BY created_at`, kbID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var result []models.KBDocument
|
||||
for rows.Next() {
|
||||
var doc models.KBDocument
|
||||
var errMsg sql.NullString
|
||||
err := rows.Scan(&doc.ID, &doc.KBID, &doc.Filename, &doc.ContentType,
|
||||
&doc.SizeBytes, &doc.StorageKey, &doc.ChunkCount, &doc.Status,
|
||||
&errMsg, &doc.UploadedBy, &doc.CreatedAt, &doc.UpdatedAt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
doc.Error = NullableStringPtr(errMsg)
|
||||
result = append(result, doc)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) UpdateDocumentStatus(ctx context.Context, id string, status string, errMsg *string) error {
|
||||
_, err := DB.ExecContext(ctx, `
|
||||
UPDATE kb_documents SET status = $2, error = $3, updated_at = now()
|
||||
WHERE id = $1`, id, status, models.NullString(errMsg))
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) UpdateDocumentText(ctx context.Context, id string, text string, chunkCount int) error {
|
||||
_, err := DB.ExecContext(ctx, `
|
||||
UPDATE kb_documents SET extracted_text = $2, chunk_count = $3, updated_at = now()
|
||||
WHERE id = $1`, id, text, chunkCount)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) DeleteDocument(ctx context.Context, id string) (*models.KBDocument, error) {
|
||||
var doc models.KBDocument
|
||||
var errMsg sql.NullString
|
||||
// Return the row before deleting so caller can clean up storage.
|
||||
err := DB.QueryRowContext(ctx, `
|
||||
DELETE FROM kb_documents WHERE id = $1
|
||||
RETURNING id, kb_id, filename, storage_key, status, error`,
|
||||
id).Scan(&doc.ID, &doc.KBID, &doc.Filename, &doc.StorageKey, &doc.Status, &errMsg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
doc.Error = NullableStringPtr(errMsg)
|
||||
return &doc, nil
|
||||
}
|
||||
|
||||
// ── Chunks ───────────────────────────────────────
|
||||
|
||||
func (s *KnowledgeBaseStore) InsertChunks(ctx context.Context, chunks []models.KBChunk) error {
|
||||
if len(chunks) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Batch insert using a single multi-row INSERT.
|
||||
// Embedding vectors are inserted as text representations that pgvector parses.
|
||||
const cols = 7 // kb_id, document_id, chunk_index, content, token_count, embedding, metadata
|
||||
valueParts := make([]string, 0, len(chunks))
|
||||
args := make([]interface{}, 0, len(chunks)*cols)
|
||||
|
||||
for i, c := range chunks {
|
||||
base := i * cols
|
||||
valueParts = append(valueParts, fmt.Sprintf(
|
||||
"($%d, $%d, $%d, $%d, $%d, $%d::vector, $%d)",
|
||||
base+1, base+2, base+3, base+4, base+5, base+6, base+7,
|
||||
))
|
||||
args = append(args, c.KBID, c.DocumentID, c.ChunkIndex,
|
||||
c.Content, c.TokenCount, vectorToString(c.Embedding), ToJSON(c.Metadata))
|
||||
}
|
||||
|
||||
q := fmt.Sprintf(`INSERT INTO kb_chunks (kb_id, document_id, chunk_index, content, token_count, embedding, metadata)
|
||||
VALUES %s`, strings.Join(valueParts, ", "))
|
||||
|
||||
_, err := DB.ExecContext(ctx, q, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) DeleteChunksForDocument(ctx context.Context, documentID string) error {
|
||||
_, err := DB.ExecContext(ctx, "DELETE FROM kb_chunks WHERE document_id = $1", documentID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) SimilaritySearch(ctx context.Context, kbIDs []string, queryVec []float64, threshold float64, limit int) ([]models.KBSearchResult, error) {
|
||||
if len(kbIDs) == 0 || len(queryVec) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 5
|
||||
}
|
||||
|
||||
rows, err := DB.QueryContext(ctx, `
|
||||
SELECT c.content, c.metadata, d.filename, kb.name,
|
||||
1 - (c.embedding <=> $1::vector) AS similarity
|
||||
FROM kb_chunks c
|
||||
JOIN kb_documents d ON c.document_id = d.id
|
||||
JOIN knowledge_bases kb ON c.kb_id = kb.id
|
||||
WHERE c.kb_id = ANY($2)
|
||||
AND 1 - (c.embedding <=> $1::vector) > $3
|
||||
ORDER BY c.embedding <=> $1::vector
|
||||
LIMIT $4`,
|
||||
vectorToString(queryVec), pq.Array(kbIDs), threshold, limit,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var results []models.KBSearchResult
|
||||
for rows.Next() {
|
||||
var r models.KBSearchResult
|
||||
var metadataJSON []byte
|
||||
err := rows.Scan(&r.Content, &metadataJSON, &r.Filename, &r.KBName, &r.Similarity)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ScanJSON(metadataJSON, &r.Metadata)
|
||||
results = append(results, r)
|
||||
}
|
||||
return results, rows.Err()
|
||||
}
|
||||
|
||||
// ── Channel Links ─────────────────────────────────
|
||||
|
||||
func (s *KnowledgeBaseStore) SetChannelKBs(ctx context.Context, channelID string, kbIDs []string) error {
|
||||
tx, err := DB.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
// Clear existing links.
|
||||
_, err = tx.ExecContext(ctx, "DELETE FROM channel_knowledge_bases WHERE channel_id = $1", channelID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Insert new links.
|
||||
for _, kbID := range kbIDs {
|
||||
_, err = tx.ExecContext(ctx, `
|
||||
INSERT INTO channel_knowledge_bases (channel_id, kb_id, enabled) VALUES ($1, $2, true)`,
|
||||
channelID, kbID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) GetChannelKBs(ctx context.Context, channelID string) ([]models.ChannelKB, error) {
|
||||
rows, err := DB.QueryContext(ctx, `
|
||||
SELECT ckb.kb_id, kb.name, ckb.enabled, kb.document_count
|
||||
FROM channel_knowledge_bases ckb
|
||||
JOIN knowledge_bases kb ON ckb.kb_id = kb.id
|
||||
WHERE ckb.channel_id = $1
|
||||
ORDER BY kb.name`, channelID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var result []models.ChannelKB
|
||||
for rows.Next() {
|
||||
var ckb models.ChannelKB
|
||||
if err := rows.Scan(&ckb.KBID, &ckb.KBName, &ckb.Enabled, &ckb.DocumentCount); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result = append(result, ckb)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
func (s *KnowledgeBaseStore) GetActiveKBIDs(ctx context.Context, channelID string, userID string, teamIDs []string) ([]string, error) {
|
||||
// Return KB IDs that are: enabled on this channel AND user has access to.
|
||||
q := `
|
||||
SELECT ckb.kb_id
|
||||
FROM channel_knowledge_bases ckb
|
||||
JOIN knowledge_bases kb ON ckb.kb_id = kb.id
|
||||
WHERE ckb.channel_id = $1 AND ckb.enabled = true
|
||||
AND (
|
||||
kb.scope = 'global'
|
||||
OR (kb.scope = 'personal' AND kb.owner_id = $2)`
|
||||
|
||||
args := []interface{}{channelID, userID}
|
||||
|
||||
if len(teamIDs) > 0 {
|
||||
q += fmt.Sprintf(` OR (kb.scope = 'team' AND kb.team_id = ANY($%d))`, len(args)+1)
|
||||
args = append(args, pq.Array(teamIDs))
|
||||
}
|
||||
|
||||
q += `)`
|
||||
|
||||
rows, err := DB.QueryContext(ctx, q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var ids []string
|
||||
for rows.Next() {
|
||||
var id string
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids, rows.Err()
|
||||
}
|
||||
|
||||
// ── Stats ────────────────────────────────────────
|
||||
|
||||
func (s *KnowledgeBaseStore) UpdateStats(ctx context.Context, kbID string) error {
|
||||
_, err := DB.ExecContext(ctx, `
|
||||
UPDATE knowledge_bases SET
|
||||
document_count = (SELECT COUNT(*) FROM kb_documents WHERE kb_id = $1 AND status != 'error'),
|
||||
chunk_count = (SELECT COUNT(*) FROM kb_chunks WHERE kb_id = $1),
|
||||
total_bytes = COALESCE((SELECT SUM(size_bytes) FROM kb_documents WHERE kb_id = $1), 0),
|
||||
updated_at = now()
|
||||
WHERE id = $1`, kbID)
|
||||
return err
|
||||
}
|
||||
|
||||
// ── Scan Helpers ─────────────────────────────────
|
||||
|
||||
func scanKB(row *sql.Row) (*models.KnowledgeBase, error) {
|
||||
var kb models.KnowledgeBase
|
||||
var ownerID, teamID sql.NullString
|
||||
var embCfgJSON []byte
|
||||
|
||||
err := row.Scan(
|
||||
&kb.ID, &kb.Name, &kb.Description, &kb.Scope,
|
||||
&ownerID, &teamID, &embCfgJSON,
|
||||
&kb.DocumentCount, &kb.ChunkCount, &kb.TotalBytes,
|
||||
&kb.Status, &kb.CreatedAt, &kb.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
kb.OwnerID = NullableStringPtr(ownerID)
|
||||
kb.TeamID = NullableStringPtr(teamID)
|
||||
ScanJSON(embCfgJSON, &kb.EmbeddingConfig)
|
||||
return &kb, nil
|
||||
}
|
||||
|
||||
func queryKBs(ctx context.Context, query string, args ...interface{}) ([]models.KnowledgeBase, error) {
|
||||
rows, err := DB.QueryContext(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var result []models.KnowledgeBase
|
||||
for rows.Next() {
|
||||
var kb models.KnowledgeBase
|
||||
var ownerID, teamID sql.NullString
|
||||
var embCfgJSON []byte
|
||||
err := rows.Scan(
|
||||
&kb.ID, &kb.Name, &kb.Description, &kb.Scope,
|
||||
&ownerID, &teamID, &embCfgJSON,
|
||||
&kb.DocumentCount, &kb.ChunkCount, &kb.TotalBytes,
|
||||
&kb.Status, &kb.CreatedAt, &kb.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
kb.OwnerID = NullableStringPtr(ownerID)
|
||||
kb.TeamID = NullableStringPtr(teamID)
|
||||
ScanJSON(embCfgJSON, &kb.EmbeddingConfig)
|
||||
result = append(result, kb)
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
// vectorToString converts a float64 slice to pgvector text format: [0.1,0.2,0.3]
|
||||
func vectorToString(v []float64) string {
|
||||
if len(v) == 0 {
|
||||
return ""
|
||||
}
|
||||
parts := make([]string, len(v))
|
||||
for i, f := range v {
|
||||
parts[i] = fmt.Sprintf("%g", f)
|
||||
}
|
||||
return "[" + strings.Join(parts, ",") + "]"
|
||||
}
|
||||
@@ -27,5 +27,6 @@ func NewStores(db *sql.DB) store.Stores {
|
||||
Pricing: NewPricingStore(),
|
||||
Extensions: NewExtensionStore(),
|
||||
Attachments: NewAttachmentStore(),
|
||||
KnowledgeBases: NewKnowledgeBaseStore(),
|
||||
}
|
||||
}
|
||||
|
||||
159
server/tools/kbsearch.go
Normal file
159
server/tools/kbsearch.go
Normal file
@@ -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
|
||||
}
|
||||
@@ -4,29 +4,98 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/lib/pq"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/database"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/knowledge"
|
||||
"git.gobha.me/xcaliber/chat-switchboard/store"
|
||||
)
|
||||
|
||||
// ── Init ────────────────────────────────────
|
||||
// ── Late Registration ────────────────────────
|
||||
// Note tools use late registration (like kb_search) because
|
||||
// note_create, note_update, and note_search need the embedder for
|
||||
// vector operations. The embedder is optional — if nil, notes still
|
||||
// work but semantic search degrades to full-text search.
|
||||
|
||||
func init() {
|
||||
Register(&NoteCreateTool{})
|
||||
Register(&NoteSearchTool{})
|
||||
Register(&NoteUpdateTool{})
|
||||
Register(&NoteListTool{})
|
||||
// RegisterNoteTools registers all note tools with captured dependencies.
|
||||
// Call from main.go after stores + embedder init.
|
||||
func RegisterNoteTools(stores store.Stores, embedder *knowledge.Embedder) {
|
||||
Register(¬eCreateTool{stores: stores, embedder: embedder})
|
||||
Register(¬eSearchTool{stores: stores, embedder: embedder})
|
||||
Register(¬eUpdateTool{stores: stores, embedder: embedder})
|
||||
Register(¬eListTool{})
|
||||
}
|
||||
|
||||
// ── Shared helpers ─────────────────────────────
|
||||
|
||||
// embedNote generates a vector for a note's title+content and stores it.
|
||||
// Silently no-ops if embedder is nil or not configured.
|
||||
func embedNote(ctx context.Context, embedder *knowledge.Embedder, stores store.Stores, noteID, userID, title, content string) {
|
||||
if embedder == nil || !embedder.IsConfigured(ctx) {
|
||||
return
|
||||
}
|
||||
|
||||
text := title + "\n\n" + content
|
||||
if len(text) > 8000 {
|
||||
text = text[:8000] // limit to ~2k tokens for embedding
|
||||
}
|
||||
|
||||
// Resolve team for embedding role
|
||||
var teamID *string
|
||||
if stores.Teams != nil {
|
||||
ids, _ := stores.Teams.GetUserTeamIDs(ctx, userID)
|
||||
if len(ids) > 0 {
|
||||
teamID = &ids[0]
|
||||
}
|
||||
}
|
||||
|
||||
result, err := embedder.EmbedChunks(ctx, userID, teamID, []string{text})
|
||||
if err != nil {
|
||||
log.Printf("⚠ note embed failed for %s: %v", noteID, err)
|
||||
return
|
||||
}
|
||||
if len(result.Vectors) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Store the vector — format as pgvector literal
|
||||
vecStr := vectorToString(result.Vectors[0])
|
||||
_, err = database.DB.Exec(
|
||||
`UPDATE notes SET embedding = $1::vector WHERE id = $2`,
|
||||
vecStr, noteID,
|
||||
)
|
||||
if err != nil {
|
||||
log.Printf("⚠ note embed store failed for %s: %v", noteID, err)
|
||||
return
|
||||
}
|
||||
|
||||
// Track embedding usage
|
||||
embedder.LogUsage(ctx, userID, nil, result)
|
||||
log.Printf("📝 note %s embedded (%d tokens)", noteID, result.InputTokens)
|
||||
}
|
||||
|
||||
// vectorToString converts a float64 slice to pgvector literal format: [0.1,0.2,...]
|
||||
func vectorToString(vec []float64) string {
|
||||
parts := make([]string, len(vec))
|
||||
for i, v := range vec {
|
||||
parts[i] = fmt.Sprintf("%g", v)
|
||||
}
|
||||
return "[" + strings.Join(parts, ",") + "]"
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════
|
||||
// note_create
|
||||
// ═══════════════════════════════════════════
|
||||
|
||||
type NoteCreateTool struct{}
|
||||
type noteCreateTool struct {
|
||||
stores store.Stores
|
||||
embedder *knowledge.Embedder
|
||||
}
|
||||
|
||||
func (t *NoteCreateTool) Definition() ToolDef {
|
||||
func (t *noteCreateTool) Definition() ToolDef {
|
||||
return ToolDef{
|
||||
Name: "note_create",
|
||||
DisplayName: "Create",
|
||||
@@ -41,7 +110,7 @@ func (t *NoteCreateTool) Definition() ToolDef {
|
||||
}
|
||||
}
|
||||
|
||||
func (t *NoteCreateTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
func (t *noteCreateTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
var args struct {
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
@@ -80,6 +149,9 @@ func (t *NoteCreateTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
return "", fmt.Errorf("failed to create note: %w", err)
|
||||
}
|
||||
|
||||
// Embed async — don't block the tool response
|
||||
go embedNote(context.Background(), t.embedder, t.stores, id, execCtx.UserID, args.Title, args.Content)
|
||||
|
||||
result, _ := json.Marshal(map[string]interface{}{
|
||||
"id": id,
|
||||
"title": args.Title,
|
||||
@@ -95,25 +167,35 @@ func (t *NoteCreateTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
// note_search
|
||||
// ═══════════════════════════════════════════
|
||||
|
||||
type NoteSearchTool struct{}
|
||||
type noteSearchTool struct {
|
||||
stores store.Stores
|
||||
embedder *knowledge.Embedder
|
||||
}
|
||||
|
||||
func (t *NoteSearchTool) Definition() ToolDef {
|
||||
func (t *noteSearchTool) Definition() ToolDef {
|
||||
return ToolDef{
|
||||
Name: "note_search",
|
||||
DisplayName: "Search",
|
||||
Category: "notes",
|
||||
Description: "Search the user's notes by keyword. Returns matching notes ranked by relevance with highlighted excerpts. Use this to find previously saved information.",
|
||||
Description: "Search the user's notes by keyword or semantic meaning. " +
|
||||
"Set semantic=true for meaning-based search using AI embeddings, " +
|
||||
"or omit for fast keyword matching. Returns matching notes ranked by relevance.",
|
||||
Parameters: JSONSchema(map[string]interface{}{
|
||||
"query": Prop("string", "Search query — natural language keywords"),
|
||||
"query": Prop("string", "Search query — natural language keywords or question"),
|
||||
"limit": Prop("integer", "Maximum results to return (default 10, max 50)"),
|
||||
"semantic": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "Use semantic (vector) search instead of keyword search. Better for meaning-based queries.",
|
||||
},
|
||||
}, []string{"query"}),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *NoteSearchTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
func (t *noteSearchTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
var args struct {
|
||||
Query string `json:"query"`
|
||||
Limit int `json:"limit"`
|
||||
Query string `json:"query"`
|
||||
Limit int `json:"limit"`
|
||||
Semantic bool `json:"semantic"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(argsJSON), &args); err != nil {
|
||||
return "", fmt.Errorf("invalid arguments: %w", err)
|
||||
@@ -126,6 +208,16 @@ func (t *NoteSearchTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
args.Limit = 10
|
||||
}
|
||||
|
||||
// Semantic search via vector similarity
|
||||
if args.Semantic {
|
||||
return t.semanticSearch(ctx, execCtx, args.Query, args.Limit)
|
||||
}
|
||||
|
||||
// Default: full-text keyword search
|
||||
return t.keywordSearch(ctx, execCtx, args.Query, args.Limit)
|
||||
}
|
||||
|
||||
func (t *noteSearchTool) keywordSearch(ctx context.Context, execCtx ExecutionContext, query string, limit int) (string, error) {
|
||||
rows, err := database.DB.Query(`
|
||||
SELECT id, title, folder_path, tags, LEFT(content, 500),
|
||||
ts_rank(search_vector, plainto_tsquery('english', $2)) AS rank,
|
||||
@@ -136,26 +228,16 @@ func (t *NoteSearchTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
AND search_vector @@ plainto_tsquery('english', $2)
|
||||
ORDER BY rank DESC
|
||||
LIMIT $3
|
||||
`, execCtx.UserID, args.Query, args.Limit)
|
||||
`, execCtx.UserID, query, limit)
|
||||
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("search failed: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
type result struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
Folder string `json:"folder"`
|
||||
Tags []string `json:"tags"`
|
||||
Excerpt string `json:"excerpt"`
|
||||
Headline string `json:"headline"`
|
||||
Rank float64 `json:"rank"`
|
||||
}
|
||||
|
||||
results := make([]result, 0)
|
||||
results := make([]noteSearchResult, 0)
|
||||
for rows.Next() {
|
||||
var r result
|
||||
var r noteSearchResult
|
||||
var dbTags pq.StringArray
|
||||
if err := rows.Scan(&r.ID, &r.Title, &r.Folder, &dbTags, &r.Excerpt, &r.Rank, &r.Headline); err != nil {
|
||||
continue
|
||||
@@ -168,20 +250,108 @@ func (t *NoteSearchTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
}
|
||||
|
||||
out, _ := json.Marshal(map[string]interface{}{
|
||||
"query": args.Query,
|
||||
"query": query,
|
||||
"mode": "keyword",
|
||||
"count": len(results),
|
||||
"results": results,
|
||||
})
|
||||
return string(out), nil
|
||||
}
|
||||
|
||||
func (t *noteSearchTool) semanticSearch(ctx context.Context, execCtx ExecutionContext, query string, limit int) (string, error) {
|
||||
// Check embedder availability — fall back to keyword search if not configured
|
||||
if t.embedder == nil || !t.embedder.IsConfigured(ctx) {
|
||||
log.Printf("⚠ note semantic search: embedder not configured, falling back to keyword")
|
||||
return t.keywordSearch(ctx, execCtx, query, limit)
|
||||
}
|
||||
|
||||
// Resolve team for embedding role
|
||||
var teamID *string
|
||||
if t.stores.Teams != nil {
|
||||
ids, _ := t.stores.Teams.GetUserTeamIDs(ctx, execCtx.UserID)
|
||||
if len(ids) > 0 {
|
||||
teamID = &ids[0]
|
||||
}
|
||||
}
|
||||
|
||||
// Embed the query
|
||||
embedResult, err := t.embedder.EmbedChunks(ctx, execCtx.UserID, teamID, []string{query})
|
||||
if err != nil {
|
||||
log.Printf("⚠ note semantic search embed failed: %v — falling back to keyword", err)
|
||||
return t.keywordSearch(ctx, execCtx, query, limit)
|
||||
}
|
||||
if len(embedResult.Vectors) == 0 {
|
||||
return t.keywordSearch(ctx, execCtx, query, limit)
|
||||
}
|
||||
|
||||
// Track embedding usage
|
||||
chanPtr := &execCtx.ChannelID
|
||||
t.embedder.LogUsage(ctx, execCtx.UserID, chanPtr, embedResult)
|
||||
|
||||
queryVec := vectorToString(embedResult.Vectors[0])
|
||||
|
||||
// Vector similarity search — only notes that have embeddings
|
||||
rows, err := database.DB.Query(`
|
||||
SELECT id, title, folder_path, tags, LEFT(content, 500),
|
||||
1 - (embedding <=> $2::vector) AS similarity
|
||||
FROM notes
|
||||
WHERE user_id = $1
|
||||
AND embedding IS NOT NULL
|
||||
AND 1 - (embedding <=> $2::vector) > 0.3
|
||||
ORDER BY embedding <=> $2::vector
|
||||
LIMIT $3
|
||||
`, execCtx.UserID, queryVec, limit)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("⚠ note semantic search query failed: %v — falling back to keyword", err)
|
||||
return t.keywordSearch(ctx, execCtx, query, limit)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
results := make([]noteSearchResult, 0)
|
||||
for rows.Next() {
|
||||
var r noteSearchResult
|
||||
var dbTags pq.StringArray
|
||||
if err := rows.Scan(&r.ID, &r.Title, &r.Folder, &dbTags, &r.Excerpt, &r.Similarity); err != nil {
|
||||
continue
|
||||
}
|
||||
r.Tags = []string(dbTags)
|
||||
if r.Tags == nil {
|
||||
r.Tags = []string{}
|
||||
}
|
||||
results = append(results, r)
|
||||
}
|
||||
|
||||
out, _ := json.Marshal(map[string]interface{}{
|
||||
"query": query,
|
||||
"mode": "semantic",
|
||||
"count": len(results),
|
||||
"results": results,
|
||||
})
|
||||
return string(out), nil
|
||||
}
|
||||
|
||||
type noteSearchResult struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
Folder string `json:"folder"`
|
||||
Tags []string `json:"tags"`
|
||||
Excerpt string `json:"excerpt"`
|
||||
Headline string `json:"headline,omitempty"`
|
||||
Rank float64 `json:"rank,omitempty"`
|
||||
Similarity float64 `json:"similarity,omitempty"`
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════
|
||||
// note_update
|
||||
// ═══════════════════════════════════════════
|
||||
|
||||
type NoteUpdateTool struct{}
|
||||
type noteUpdateTool struct {
|
||||
stores store.Stores
|
||||
embedder *knowledge.Embedder
|
||||
}
|
||||
|
||||
func (t *NoteUpdateTool) Definition() ToolDef {
|
||||
func (t *noteUpdateTool) Definition() ToolDef {
|
||||
return ToolDef{
|
||||
Name: "note_update",
|
||||
DisplayName: "Update",
|
||||
@@ -197,7 +367,7 @@ func (t *NoteUpdateTool) Definition() ToolDef {
|
||||
}
|
||||
}
|
||||
|
||||
func (t *NoteUpdateTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
func (t *noteUpdateTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
var args struct {
|
||||
NoteID string `json:"note_id"`
|
||||
Title *string `json:"title"`
|
||||
@@ -251,14 +421,17 @@ func (t *NoteUpdateTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
queryArgs = append(queryArgs, args.NoteID, execCtx.UserID)
|
||||
query := "UPDATE notes SET " + strings.Join(setClauses, ", ") +
|
||||
fmt.Sprintf(" WHERE id = $%d AND user_id = $%d", argIdx, argIdx+1) +
|
||||
" RETURNING id, title, updated_at::text"
|
||||
" RETURNING id, title, content, updated_at::text"
|
||||
|
||||
var id, title, updatedAt string
|
||||
err := database.DB.QueryRow(query, queryArgs...).Scan(&id, &title, &updatedAt)
|
||||
var id, title, content, updatedAt string
|
||||
err := database.DB.QueryRow(query, queryArgs...).Scan(&id, &title, &content, &updatedAt)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("note not found or update failed")
|
||||
}
|
||||
|
||||
// Re-embed async with updated content
|
||||
go embedNote(context.Background(), t.embedder, t.stores, id, execCtx.UserID, title, content)
|
||||
|
||||
result, _ := json.Marshal(map[string]interface{}{
|
||||
"id": id,
|
||||
"title": title,
|
||||
@@ -272,9 +445,9 @@ func (t *NoteUpdateTool) Execute(ctx context.Context, execCtx ExecutionContext,
|
||||
// note_list
|
||||
// ═══════════════════════════════════════════
|
||||
|
||||
type NoteListTool struct{}
|
||||
type noteListTool struct{}
|
||||
|
||||
func (t *NoteListTool) Definition() ToolDef {
|
||||
func (t *noteListTool) Definition() ToolDef {
|
||||
return ToolDef{
|
||||
Name: "note_list",
|
||||
DisplayName: "List",
|
||||
@@ -288,7 +461,7 @@ func (t *NoteListTool) Definition() ToolDef {
|
||||
}
|
||||
}
|
||||
|
||||
func (t *NoteListTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
func (t *noteListTool) Execute(ctx context.Context, execCtx ExecutionContext, argsJSON string) (string, error) {
|
||||
var args struct {
|
||||
Folder string `json:"folder"`
|
||||
Tag string `json:"tag"`
|
||||
|
||||
@@ -3,13 +3,23 @@ package tools
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/store"
|
||||
)
|
||||
|
||||
// TestMain registers late-registered tools (notes) with nil deps so
|
||||
// that Definition()-level tests work. Execute() tests need real deps.
|
||||
func TestMain(m *testing.M) {
|
||||
RegisterNoteTools(store.Stores{}, nil)
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
// ── Registry Tests ──────────────────────────
|
||||
|
||||
func TestRegistryHasTools(t *testing.T) {
|
||||
// init() in notes.go registers 4 tools
|
||||
// TestMain registers note tools via RegisterNoteTools
|
||||
if !HasTools() {
|
||||
t.Fatal("Expected HasTools() == true after init registration")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user