archivdms/internal/storage/metadata_suggestions_naivebayes.go

package storage

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

	"archivdms/internal/classifier"
)

// GenerateNaiveBayesSuggestions builds a metadata suggestion for a document
// using the trained Naive-Bayes model (internal/classifier) instead of the
// fuzzy-name heuristic or an LLM. It classifies the document's title+OCR text
// against the tenant's trained document_types / correspondents / tags models,
// maps the predicted class IDs back to taxonomy entities, drops entities that
// are already assigned, and persists the result as a metadata_suggestions row
// with provider='naive_bayes' — in the SAME SuggestionPayload schema the other
// providers produce, so the API/frontend are unchanged.
//
// A kind whose model is untrained (or below the per-class data threshold) simply
// yields no candidates for that kind — not an error. Any real failure (DB error,
// classifier error) is returned as-is: there is NO silent fallback to the
// heuristic provider (GoBD-Nachvollziehbarkeit — the caller reports which
// provider produced or failed the run). requestedBy may be nil for
// non-interactive callers.
func (s *Store) GenerateNaiveBayesSuggestions(ctx context.Context, documentID, tenantID int64, requestedBy *int64) (*MetadataSuggestion, error) {
	doc, err := s.GetDocument(ctx, documentID, tenantID)
	if err != nil {
		return nil, err // ErrDocumentNotFound propagates
	}

	text := doc.Title
	if doc.OCRText != "" {
		text = doc.Title + "\n" + doc.OCRText
	}

	clf := classifier.New(s.db)

	// Entities already assigned are excluded from suggestions, matching the
	// other providers' behaviour.
	assignedTags := map[int64]bool{}
	docTags, err := s.ListDocumentTags(ctx, documentID, tenantID)
	if err != nil {
		return nil, err
	}
	for _, t := range docTags {
		assignedTags[t.ID] = true
	}

	docTypeCands, err := s.naiveBayesCandidates(ctx, clf, "document_types", tenantID, text, func(id int64) bool {
		return doc.DocTypeID != nil && *doc.DocTypeID == id
	})
	if err != nil {
		return nil, err
	}
	corrCands, err := s.naiveBayesCandidates(ctx, clf, "correspondents", tenantID, text, func(id int64) bool {
		return doc.CorrespondentID != nil && *doc.CorrespondentID == id
	})
	if err != nil {
		return nil, err
	}
	tagCands, err := s.naiveBayesCandidates(ctx, clf, "tags", tenantID, text, func(id int64) bool {
		return assignedTags[id]
	})
	if err != nil {
		return nil, err
	}

	payload := SuggestionPayload{
		DocTypeCandidates:       docTypeCands,
		CorrespondentCandidates: corrCands,
		TagCandidates:           tagCands,
	}
	// The Naive-Bayes model does not propose a title (it classifies against
	// existing entities only); Title stays nil.

	raw, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("storage: marshal naive_bayes suggestion payload: %w", err)
	}

	row := s.db.QueryRow(ctx, `
		INSERT INTO metadata_suggestions (tenant_id, document_id, provider, requested_by, suggestion)
		VALUES ($1, $2, 'naive_bayes', $3, $4)
		RETURNING `+metadataSuggestionCols,
		tenantID, documentID, requestedBy, raw)
	m, err := scanMetadataSuggestion(row)
	if err != nil {
		return nil, fmt.Errorf("storage: insert naive_bayes metadata suggestion: %w", err)
	}
	return m, nil
}

// naiveBayesCandidates runs the classifier for one kind and maps predicted class
// IDs back to SuggestionCandidate (resolving the entity name from the taxonomy),
// dropping excluded (already-assigned) entities and any predicted ID that no
// longer exists as a live entity. Result is always non-nil.
func (s *Store) naiveBayesCandidates(ctx context.Context, clf *classifier.Classifier, kind string, tenantID int64, text string, excluded func(id int64) bool) ([]SuggestionCandidate, error) {
	preds, err := clf.Predict(ctx, tenantID, kind, text)
	if err != nil {
		return nil, fmt.Errorf("storage: naive_bayes predict %s: %w", kind, err)
	}
	if len(preds) == 0 {
		return make([]SuggestionCandidate, 0), nil
	}

	entities, err := s.ListTaxonomyEntities(ctx, kind, tenantID)
	if err != nil {
		return nil, err
	}
	names := make(map[int64]string, len(entities))
	for _, e := range entities {
		names[e.ID] = e.Name
	}

	out := make([]SuggestionCandidate, 0, len(preds))
	for _, p := range preds {
		if excluded(p.EntityID) {
			continue
		}
		name, ok := names[p.EntityID]
		if !ok {
			continue // predicted a class whose entity was deleted since training
		}
		out = append(out, SuggestionCandidate{ID: p.EntityID, Name: name, Score: p.Score, Explanation: p.TopTokens})
	}
	return out, nil
}