Changeset 0.14.0 (#67)
This commit is contained in:
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
|
||||
}
|
||||
Reference in New Issue
Block a user