patrick
9a24ea29e1
Git-Repository für bestehenden archivdms-Code initialisiert, Branch-/Commit-Konvention (feature/<ticket>-<slug>-Branches, Ticket-Prefix in Commit-Nachricht) etabliert.
60 lines
2.1 KiB
Go
60 lines
2.1 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"archivdms/internal/classifier"
|
|
)
|
|
|
|
// MLClassifierKinds is the fixed set of taxonomy kinds the Naive-Bayes
|
|
// classifier is trained/predicted for, in a stable order (used by the retrain
|
|
// CLI so runs are deterministic).
|
|
var MLClassifierKinds = []string{
|
|
classifier.KindDocumentTypes,
|
|
classifier.KindCorrespondents,
|
|
classifier.KindTags,
|
|
}
|
|
|
|
// TrainClassifier rebuilds the Naive-Bayes model for one tenant and one kind and
|
|
// returns the number of training documents used. Thin wrapper around
|
|
// classifier.Train that keeps the Store's *pgxpool.Pool encapsulated (the
|
|
// classifier package must not import storage). See classifier.Train for
|
|
// semantics (full rebuild, MinDocsPerClass threshold, no error on empty data).
|
|
func (s *Store) TrainClassifier(ctx context.Context, tenantID int64, kind string) (int, error) {
|
|
return classifier.New(s.db).Train(ctx, tenantID, kind)
|
|
}
|
|
|
|
// StartMLRun inserts a fresh ml_classifier_runs row in status 'running' for a
|
|
// tenant and returns its id. The retrain CLI creates one run per tenant that
|
|
// spans all kinds, then finalises it with FinishMLRun.
|
|
func (s *Store) StartMLRun(ctx context.Context, tenantID int64) (int64, error) {
|
|
var id int64
|
|
if err := s.db.QueryRow(ctx,
|
|
`INSERT INTO ml_classifier_runs (tenant_id, status) VALUES ($1, 'running') RETURNING id`,
|
|
tenantID).Scan(&id); err != nil {
|
|
return 0, fmt.Errorf("storage: start ml_classifier run: %w", err)
|
|
}
|
|
return id, nil
|
|
}
|
|
|
|
// FinishMLRun finalises an ml_classifier_runs row: sets completed_at=now(),
|
|
// doc_count, the terminal status ('completed', 'failed' or
|
|
// 'skipped_insufficient_data') and an optional error message. errMsg "" stores
|
|
// SQL NULL.
|
|
func (s *Store) FinishMLRun(ctx context.Context, runID int64, docCount int, status, errMsg string) error {
|
|
var errArg any
|
|
if errMsg != "" {
|
|
errArg = errMsg
|
|
}
|
|
_, err := s.db.Exec(ctx,
|
|
`UPDATE ml_classifier_runs
|
|
SET completed_at = now(), doc_count = $2, status = $3, error = $4
|
|
WHERE id = $1`,
|
|
runID, docCount, status, errArg)
|
|
if err != nil {
|
|
return fmt.Errorf("storage: finish ml_classifier run: %w", err)
|
|
}
|
|
return nil
|
|
}
|