feat(captain): queue document crawl jobs
This commit is contained in:
@@ -582,6 +582,9 @@ func Bootstrap(env string) (*App, error) {
|
||||
captainAssistantService := service.NewCaptainAssistantService(captainAssistantRepo, captainInboxRepo, captainDocumentRepo, captainAssistantResponseRepo, llmProvider)
|
||||
captainDocumentService := service.NewCaptainDocumentService(captainDocumentRepo, llmProvider, captainAssistantRepo)
|
||||
captainDocumentService.SetWorkerPool(workerPool)
|
||||
if _, err := service.EnqueueCaptainDocumentScheduleSyncs(context.Background(), workerPool, time.Now()); err != nil {
|
||||
applogger.L().Warnf("failed to enqueue initial Captain document sync scheduler: %v", err)
|
||||
}
|
||||
captainScenarioService := service.NewCaptainScenarioService(captainScenarioRepo, captainAssistantRepo)
|
||||
captainCustomToolService := service.NewCaptainCustomToolService(captainCustomToolRepo)
|
||||
copilotService := service.NewCopilotService(copilotThreadRepo, copilotMessageRepo, copilotSuggestionRepo, llmProvider, captainAssistantRepo)
|
||||
|
||||
@@ -3,6 +3,7 @@ package repository
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gochat/gochat/internal/model"
|
||||
"gorm.io/gorm"
|
||||
@@ -123,6 +124,37 @@ func (r *CaptainDocumentRepo) FindByExternalLink(ctx context.Context, assistantI
|
||||
return &doc, nil
|
||||
}
|
||||
|
||||
func (r *CaptainDocumentRepo) ListDueForAutoSync(ctx context.Context, now time.Time, interval, staleTimeout time.Duration, limit int) ([]model.CaptainDocument, error) {
|
||||
if interval <= 0 {
|
||||
interval = 24 * time.Hour
|
||||
}
|
||||
if staleTimeout <= 0 {
|
||||
staleTimeout = 10 * time.Minute
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 1000
|
||||
}
|
||||
syncDueBefore := now.Add(-interval / 2).Unix()
|
||||
staleCutoff := now.Add(-staleTimeout).Unix()
|
||||
|
||||
var docs []model.CaptainDocument
|
||||
err := r.db.WithContext(ctx).
|
||||
Model(&model.CaptainDocument{}).
|
||||
Where("status = ?", model.DocumentStatusCompleted).
|
||||
Where(
|
||||
"(sync_status = ? AND (last_synced_at IS NULL OR last_synced_at < ?)) OR "+
|
||||
"(sync_status = ? AND (last_sync_attempted_at IS NULL OR last_sync_attempted_at < ?)) OR "+
|
||||
"(sync_status = ? AND last_sync_attempted_at < ?)",
|
||||
model.DocumentSyncStatusSynced, syncDueBefore,
|
||||
model.DocumentSyncStatusFailed, syncDueBefore,
|
||||
model.DocumentSyncStatusPending, staleCutoff,
|
||||
).
|
||||
Order("last_sync_attempted_at ASC, id ASC").
|
||||
Limit(limit).
|
||||
Find(&docs).Error
|
||||
return docs, err
|
||||
}
|
||||
|
||||
// FindByStatus retrieves documents matching a specific status.
|
||||
func (r *CaptainDocumentRepo) FindByStatus(ctx context.Context, assistantID uint, status model.DocumentStatus, offset, limit int) ([]model.CaptainDocument, int64, error) {
|
||||
var docs []model.CaptainDocument
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -14,6 +15,7 @@ import (
|
||||
"github.com/gochat/gochat/internal/repository"
|
||||
"github.com/gochat/gochat/internal/worker"
|
||||
applogger "github.com/gochat/gochat/pkg/logger"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// CaptainDocumentService implements business logic for CaptainDocument operations.
|
||||
@@ -23,6 +25,8 @@ type CaptainDocumentService struct {
|
||||
assistantRepo *repository.CaptainAssistantRepo
|
||||
llmProvider llm.Provider
|
||||
syncBackend CaptainDocumentSyncBackend
|
||||
crawlBackend CaptainDocumentCrawlBackend
|
||||
pageParser CaptainDocumentPageParserBackend
|
||||
worker *worker.WorkerPool
|
||||
}
|
||||
|
||||
@@ -36,6 +40,19 @@ type CaptainDocumentSyncResult struct {
|
||||
ErrorCode string
|
||||
}
|
||||
|
||||
type CaptainDocumentCrawlBackend interface {
|
||||
CrawlCaptainDocument(ctx context.Context, doc *model.CaptainDocument) (*CaptainDocumentCrawlResult, error)
|
||||
}
|
||||
|
||||
type CaptainDocumentCrawlResult struct {
|
||||
PageLinks []string
|
||||
ErrorCode string
|
||||
}
|
||||
|
||||
type CaptainDocumentPageParserBackend interface {
|
||||
ParseCaptainDocumentPage(ctx context.Context, pageLink string) (*CaptainDocumentSyncResult, error)
|
||||
}
|
||||
|
||||
// NewCaptainDocumentService creates a new CaptainDocumentService.
|
||||
func NewCaptainDocumentService(
|
||||
documentRepo *repository.CaptainDocumentRepo,
|
||||
@@ -56,6 +73,14 @@ func (s *CaptainDocumentService) SetSyncBackend(syncBackend CaptainDocumentSyncB
|
||||
s.syncBackend = syncBackend
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetCrawlBackend(crawlBackend CaptainDocumentCrawlBackend) {
|
||||
s.crawlBackend = crawlBackend
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetPageParserBackend(pageParser CaptainDocumentPageParserBackend) {
|
||||
s.pageParser = pageParser
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetWorkerPool(wp *worker.WorkerPool) {
|
||||
s.worker = wp
|
||||
RegisterCaptainDocumentJobs(wp, s)
|
||||
@@ -108,8 +133,14 @@ func (s *CaptainDocumentService) Create(ctx context.Context, assistantID, accoun
|
||||
return nil, fmt.Errorf("create document: %w", err)
|
||||
}
|
||||
if created, err := s.documentRepo.GetByAccountAndID(ctx, accountID, doc.ID); err == nil {
|
||||
if enqueueErr := s.enqueueDocumentCrawl(ctx, accountID, created.ID); enqueueErr != nil {
|
||||
return nil, enqueueErr
|
||||
}
|
||||
return created, nil
|
||||
}
|
||||
if enqueueErr := s.enqueueDocumentCrawl(ctx, accountID, doc.ID); enqueueErr != nil {
|
||||
return nil, enqueueErr
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
@@ -237,6 +268,166 @@ func (s *CaptainDocumentService) RequestSyncDocumentByAccount(ctx context.Contex
|
||||
return doc, nil
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) RequestCrawlDocumentByAccount(ctx context.Context, accountID, id uint) (*model.CaptainDocument, error) {
|
||||
doc, err := s.MarkSyncing(ctx, accountID, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
doc.Status = model.DocumentStatusInProgress
|
||||
if err := s.documentRepo.Update(ctx, doc); err != nil {
|
||||
return nil, fmt.Errorf("mark document crawling: %w", err)
|
||||
}
|
||||
if err := s.enqueueDocumentCrawl(ctx, accountID, id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.documentRepo.GetByAccountAndID(ctx, accountID, id)
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) enqueueDocumentCrawl(ctx context.Context, accountID, id uint) error {
|
||||
if s.worker == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := s.worker.Enqueue(ctx, TaskTypeCaptainDocumentCrawl, captainDocumentCrawlJob{AccountID: accountID, DocumentID: id},
|
||||
worker.WithQueue("low"),
|
||||
worker.WithMaxAttempts(3),
|
||||
worker.WithIdempotencyKey(fmt.Sprintf("captain:document_crawl:%d:%d", accountID, id)),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) ScheduleDueDocumentSyncs(ctx context.Context, now time.Time) (int, error) {
|
||||
if s.worker == nil {
|
||||
return 0, nil
|
||||
}
|
||||
docs, err := s.documentRepo.ListDueForAutoSync(ctx, now, 24*time.Hour, 10*time.Minute, 1000)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("list due captain documents: %w", err)
|
||||
}
|
||||
enqueued := 0
|
||||
for _, doc := range docs {
|
||||
if !s.captainDocumentAutoSyncEnabled(ctx, doc.AccountID) {
|
||||
continue
|
||||
}
|
||||
if _, err := s.worker.Enqueue(ctx, TaskTypeCaptainDocumentSync, captainDocumentSyncJob{AccountID: doc.AccountID, DocumentID: doc.ID},
|
||||
worker.WithQueue("purgable"),
|
||||
worker.WithMaxAttempts(3),
|
||||
worker.WithIdempotencyKey(fmt.Sprintf("captain:document_sync:auto:%d:%d:%d", doc.AccountID, doc.ID, now.UTC().Truncate(24*time.Hour).Unix())),
|
||||
); err != nil {
|
||||
return enqueued, err
|
||||
}
|
||||
enqueued++
|
||||
}
|
||||
return enqueued, nil
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) captainDocumentAutoSyncEnabled(ctx context.Context, accountID uint) bool {
|
||||
var account model.Account
|
||||
if err := s.documentRepo.DB().WithContext(ctx).Select("feature_flags").First(&account, accountID).Error; err != nil {
|
||||
return false
|
||||
}
|
||||
return featureFlagStringEnabled(account.FeatureFlags, "captain_document_auto_sync")
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) CrawlDocumentByAccount(ctx context.Context, accountID, id uint) (*model.CaptainDocument, error) {
|
||||
doc, err := s.documentRepo.GetByAccountAndID(ctx, accountID, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("document not found: %w", err)
|
||||
}
|
||||
if s.crawlBackend == nil {
|
||||
return s.markDocumentSyncFailed(ctx, accountID, id, "crawl_disabled")
|
||||
}
|
||||
result, err := s.crawlBackend.CrawlCaptainDocument(ctx, doc)
|
||||
if err != nil {
|
||||
updated, markErr := s.markDocumentSyncFailed(ctx, accountID, id, "crawl_error")
|
||||
if markErr != nil {
|
||||
return nil, markErr
|
||||
}
|
||||
return updated, fmt.Errorf("crawl document: %w", err)
|
||||
}
|
||||
if result == nil {
|
||||
return s.markDocumentSyncFailed(ctx, accountID, id, "crawl_error")
|
||||
}
|
||||
if result.ErrorCode != "" {
|
||||
return s.markDocumentSyncFailed(ctx, accountID, id, result.ErrorCode)
|
||||
}
|
||||
if s.worker == nil {
|
||||
return doc, fmt.Errorf("worker pool required for document crawl parser fan-out")
|
||||
}
|
||||
links := normalizedUniqueLinks(append(result.PageLinks, doc.ExternalLink))
|
||||
for _, link := range links {
|
||||
if _, err := s.worker.Enqueue(ctx, TaskTypeCaptainDocumentPageCrawlParse, captainDocumentPageCrawlParseJob{AccountID: accountID, AssistantID: doc.AssistantID, PageLink: link},
|
||||
worker.WithQueue("low"),
|
||||
worker.WithMaxAttempts(3),
|
||||
worker.WithIdempotencyKey(fmt.Sprintf("captain:document_page_crawl_parse:%d:%d:%s", accountID, doc.AssistantID, link)),
|
||||
); err != nil {
|
||||
return doc, err
|
||||
}
|
||||
}
|
||||
return s.documentRepo.GetByAccountAndID(ctx, accountID, id)
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) ParseCrawledPage(ctx context.Context, accountID, assistantID uint, pageLink string) (*model.CaptainDocument, error) {
|
||||
pageLink = normalizeCaptainDocumentLink(pageLink)
|
||||
if accountID == 0 || assistantID == 0 || pageLink == "" {
|
||||
return nil, fmt.Errorf("invalid captain page crawl payload: account_id=%d assistant_id=%d page_link=%q", accountID, assistantID, pageLink)
|
||||
}
|
||||
if s.assistantRepo != nil {
|
||||
if _, err := s.assistantRepo.GetByAccountAndID(ctx, accountID, assistantID); err != nil {
|
||||
return nil, fmt.Errorf("assistant not found: %w", err)
|
||||
}
|
||||
}
|
||||
doc, err := s.documentRepo.FindByExternalLink(ctx, assistantID, pageLink)
|
||||
if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, err
|
||||
}
|
||||
if s.pageParser == nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, fmt.Errorf("page crawl parser disabled")
|
||||
}
|
||||
return s.markDocumentSyncFailed(ctx, accountID, doc.ID, "crawl_disabled")
|
||||
}
|
||||
result, parseErr := s.pageParser.ParseCaptainDocumentPage(ctx, pageLink)
|
||||
if parseErr != nil {
|
||||
if err == nil {
|
||||
_, _ = s.markDocumentSyncFailed(ctx, accountID, doc.ID, "fetch_failed")
|
||||
}
|
||||
return nil, fmt.Errorf("parse crawled page: %w", parseErr)
|
||||
}
|
||||
if result == nil || result.ErrorCode != "" || strings.TrimSpace(result.Content) == "" {
|
||||
code := "content_empty"
|
||||
if result != nil && result.ErrorCode != "" {
|
||||
code = result.ErrorCode
|
||||
}
|
||||
if err == nil {
|
||||
return s.markDocumentSyncFailed(ctx, accountID, doc.ID, code)
|
||||
}
|
||||
return nil, fmt.Errorf("parse crawled page failed: %s", code)
|
||||
}
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
doc = &model.CaptainDocument{AccountID: accountID, AssistantID: assistantID, ExternalLink: pageLink}
|
||||
}
|
||||
doc.Name = strings.TrimSpace(result.Title)
|
||||
if doc.Name == "" {
|
||||
doc.Name = pageLink
|
||||
}
|
||||
doc.Content = strings.TrimSpace(result.Content)
|
||||
doc.ContentFingerprint = computeFingerprint(doc.Content)
|
||||
doc.Status = model.DocumentStatusCompleted
|
||||
doc.SyncStatus = model.DocumentSyncStatusSynced
|
||||
doc.LastSyncErrorCode = ""
|
||||
now := time.Now().Unix()
|
||||
doc.LastSyncedAt = &now
|
||||
doc.LastSyncAttemptedAt = &now
|
||||
if doc.ID == 0 {
|
||||
if err := s.documentRepo.Create(ctx, doc); err != nil {
|
||||
return nil, fmt.Errorf("create crawled document: %w", err)
|
||||
}
|
||||
} else if err := s.documentRepo.Update(ctx, doc); err != nil {
|
||||
return nil, fmt.Errorf("update crawled document: %w", err)
|
||||
}
|
||||
return s.documentRepo.GetByAccountAndID(ctx, accountID, doc.ID)
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SyncDocumentByAccount(ctx context.Context, accountID, id uint) (*model.CaptainDocument, error) {
|
||||
doc, err := s.documentRepo.GetByAccountAndID(ctx, accountID, id)
|
||||
if err != nil {
|
||||
@@ -452,3 +643,24 @@ func computeFingerprint(content string) string {
|
||||
h.Write([]byte(normalized))
|
||||
return fmt.Sprintf("%x", h.Sum(nil))
|
||||
}
|
||||
|
||||
func normalizedUniqueLinks(rawLinks []string) []string {
|
||||
seen := make(map[string]struct{}, len(rawLinks))
|
||||
links := make([]string, 0, len(rawLinks))
|
||||
for _, rawLink := range rawLinks {
|
||||
link := normalizeCaptainDocumentLink(rawLink)
|
||||
if link == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[link]; ok {
|
||||
continue
|
||||
}
|
||||
seen[link] = struct{}{}
|
||||
links = append(links, link)
|
||||
}
|
||||
return links
|
||||
}
|
||||
|
||||
func normalizeCaptainDocumentLink(rawLink string) string {
|
||||
return strings.TrimRight(strings.TrimSpace(rawLink), "/")
|
||||
}
|
||||
|
||||
@@ -155,6 +155,117 @@ func TestCaptainDocumentService_DocumentSyncJobRetriesMissingDocument(t *testing
|
||||
assert.NotEmpty(t, job.LastError)
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_RequestCrawlQueuesParserJobs(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
now := time.Date(2026, 6, 5, 23, 30, 0, 0, time.UTC)
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithNow(func() time.Time { return now }))
|
||||
svc.SetWorkerPool(wp)
|
||||
svc.SetCrawlBackend(&captainDocumentFakeCrawlBackend{links: []string{"https://example.com/help", "https://example.com/faq/"}})
|
||||
svc.SetPageParserBackend(&captainDocumentFakePageParserBackend{pages: map[string]*CaptainDocumentSyncResult{
|
||||
"https://example.com/help": {Title: "Fresh Help", Content: "fresh help content"},
|
||||
"https://example.com/faq": {Title: "FAQ", Content: "answer one"},
|
||||
}})
|
||||
|
||||
queued, err := svc.RequestCrawlDocumentByAccount(context.Background(), account.ID, doc.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, model.DocumentStatusInProgress, queued.Status)
|
||||
assert.Equal(t, model.DocumentSyncStatusPending, queued.SyncStatus)
|
||||
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var parserJobCount int64
|
||||
require.NoError(t, db.Model(&model.BackgroundJob{}).Where("job_type = ? AND queue = ?", TaskTypeCaptainDocumentPageCrawlParse, "low").Count(&parserJobCount).Error)
|
||||
assert.Equal(t, int64(2), parserJobCount)
|
||||
|
||||
processed, err = wp.ProcessOne(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, processed)
|
||||
processed, err = wp.ProcessOne(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var updated model.CaptainDocument
|
||||
require.NoError(t, db.First(&updated, doc.ID).Error)
|
||||
assert.Equal(t, "Fresh Help", updated.Name)
|
||||
assert.Equal(t, "fresh help content", updated.Content)
|
||||
assert.Equal(t, model.DocumentStatusCompleted, updated.Status)
|
||||
assert.Equal(t, model.DocumentSyncStatusSynced, updated.SyncStatus)
|
||||
|
||||
var created model.CaptainDocument
|
||||
require.NoError(t, db.Where("assistant_id = ? AND external_link = ?", doc.AssistantID, "https://example.com/faq").First(&created).Error)
|
||||
assert.Equal(t, account.ID, created.AccountID)
|
||||
assert.Equal(t, "FAQ", created.Name)
|
||||
assert.Equal(t, "answer one", created.Content)
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_CrawlJobDisabledMarksFailed(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
wp := worker.NewWorkerPool(db)
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
_, err := svc.RequestCrawlDocumentByAccount(context.Background(), account.ID, doc.ID)
|
||||
require.NoError(t, err)
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var updated model.CaptainDocument
|
||||
require.NoError(t, db.First(&updated, doc.ID).Error)
|
||||
assert.Equal(t, model.DocumentSyncStatusFailed, updated.SyncStatus)
|
||||
assert.Equal(t, "crawl_disabled", updated.LastSyncErrorCode)
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_ScheduleDueDocumentSyncsQueuesDueDocs(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
account.FeatureFlags = `{"captain_document_auto_sync":true}`
|
||||
require.NoError(t, db.Save(account).Error)
|
||||
now := time.Date(2026, 6, 6, 0, 0, 0, 0, time.UTC)
|
||||
oldSync := now.Add(-13 * time.Hour).Unix()
|
||||
doc.LastSyncedAt = &oldSync
|
||||
doc.LastSyncAttemptedAt = &oldSync
|
||||
doc.Status = model.DocumentStatusCompleted
|
||||
doc.SyncStatus = model.DocumentSyncStatusSynced
|
||||
require.NoError(t, db.Save(doc).Error)
|
||||
freshSync := now.Add(-time.Hour).Unix()
|
||||
fresh := &model.CaptainDocument{AccountID: account.ID, AssistantID: doc.AssistantID, Name: "Fresh", ExternalLink: "https://example.com/fresh", Status: model.DocumentStatusCompleted, SyncStatus: model.DocumentSyncStatusSynced, LastSyncedAt: &freshSync, LastSyncAttemptedAt: &freshSync}
|
||||
require.NoError(t, db.Create(fresh).Error)
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithNow(func() time.Time { return now }))
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
enqueued, err := svc.ScheduleDueDocumentSyncs(context.Background(), now)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, enqueued)
|
||||
|
||||
var jobs []model.BackgroundJob
|
||||
require.NoError(t, db.Where("job_type = ?", TaskTypeCaptainDocumentSync).Find(&jobs).Error)
|
||||
require.Len(t, jobs, 1)
|
||||
assert.Equal(t, "purgable", jobs[0].Queue)
|
||||
assert.Contains(t, string(jobs[0].Payload), fmt.Sprintf("\"document_id\":%d", doc.ID))
|
||||
}
|
||||
|
||||
func TestEnqueueCaptainDocumentScheduleSyncsUsesDailyIdempotency(t *testing.T) {
|
||||
db, _ := setupCaptainDocumentServiceTest(t)
|
||||
now := time.Date(2026, 6, 6, 1, 0, 0, 0, time.UTC)
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithNow(func() time.Time { return now }))
|
||||
|
||||
first, err := EnqueueCaptainDocumentScheduleSyncs(context.Background(), wp, now)
|
||||
require.NoError(t, err)
|
||||
second, err := EnqueueCaptainDocumentScheduleSyncs(context.Background(), wp, now.Add(time.Hour))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, first.ID, second.ID)
|
||||
assert.Equal(t, TaskTypeCaptainDocumentScheduleSyncs, first.JobType)
|
||||
assert.Equal(t, "scheduled_jobs", first.Queue)
|
||||
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&model.BackgroundJob{}).Where("job_type = ?", TaskTypeCaptainDocumentScheduleSyncs).Count(&count).Error)
|
||||
assert.Equal(t, int64(1), count)
|
||||
}
|
||||
|
||||
type captainDocumentFakeSyncBackend struct {
|
||||
result *CaptainDocumentSyncResult
|
||||
err error
|
||||
@@ -168,3 +279,29 @@ func (b *captainDocumentFakeSyncBackend) SyncCaptainDocument(ctx context.Context
|
||||
}
|
||||
return b.result, nil
|
||||
}
|
||||
|
||||
type captainDocumentFakeCrawlBackend struct {
|
||||
links []string
|
||||
err error
|
||||
documentID uint
|
||||
}
|
||||
|
||||
func (b *captainDocumentFakeCrawlBackend) CrawlCaptainDocument(ctx context.Context, doc *model.CaptainDocument) (*CaptainDocumentCrawlResult, error) {
|
||||
b.documentID = doc.ID
|
||||
if b.err != nil {
|
||||
return nil, b.err
|
||||
}
|
||||
return &CaptainDocumentCrawlResult{PageLinks: b.links}, nil
|
||||
}
|
||||
|
||||
type captainDocumentFakePageParserBackend struct {
|
||||
pages map[string]*CaptainDocumentSyncResult
|
||||
err error
|
||||
}
|
||||
|
||||
func (b *captainDocumentFakePageParserBackend) ParseCaptainDocumentPage(ctx context.Context, pageLink string) (*CaptainDocumentSyncResult, error) {
|
||||
if b.err != nil {
|
||||
return nil, b.err
|
||||
}
|
||||
return b.pages[pageLink], nil
|
||||
}
|
||||
|
||||
@@ -5,18 +5,41 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gochat/gochat/internal/model"
|
||||
"github.com/gochat/gochat/internal/worker"
|
||||
)
|
||||
|
||||
const TaskTypeCaptainDocumentSync = "captain:document_sync"
|
||||
const (
|
||||
TaskTypeCaptainDocumentSync = "captain:document_sync"
|
||||
TaskTypeCaptainDocumentCrawl = "captain:document_crawl"
|
||||
TaskTypeCaptainDocumentPageCrawlParse = "captain:document_page_crawl_parse"
|
||||
TaskTypeCaptainDocumentScheduleSyncs = "captain:documents_schedule_syncs"
|
||||
)
|
||||
|
||||
const captainDocumentScheduleInterval = 24 * time.Hour
|
||||
|
||||
type captainDocumentSyncJob struct {
|
||||
AccountID uint `json:"account_id"`
|
||||
DocumentID uint `json:"document_id"`
|
||||
}
|
||||
|
||||
type captainDocumentCrawlJob struct {
|
||||
AccountID uint `json:"account_id"`
|
||||
DocumentID uint `json:"document_id"`
|
||||
}
|
||||
|
||||
type captainDocumentPageCrawlParseJob struct {
|
||||
AccountID uint `json:"account_id"`
|
||||
AssistantID uint `json:"assistant_id"`
|
||||
PageLink string `json:"page_link"`
|
||||
}
|
||||
|
||||
type captainDocumentScheduleSyncsJob struct {
|
||||
PlanName string `json:"plan_name,omitempty"`
|
||||
}
|
||||
|
||||
var captainDocumentRegistrations sync.Map
|
||||
|
||||
// RegisterCaptainDocumentJobs wires Captain::Documents::PerformSyncJob into
|
||||
@@ -30,6 +53,29 @@ func RegisterCaptainDocumentJobs(wp *worker.WorkerPool, svc *CaptainDocumentServ
|
||||
return
|
||||
}
|
||||
wp.Register(TaskTypeCaptainDocumentSync, svc.performDocumentSyncJob)
|
||||
wp.Register(TaskTypeCaptainDocumentCrawl, svc.performDocumentCrawlJob)
|
||||
wp.Register(TaskTypeCaptainDocumentPageCrawlParse, svc.performDocumentPageCrawlParseJob)
|
||||
wp.Register(TaskTypeCaptainDocumentScheduleSyncs, svc.performDocumentScheduleSyncsJob)
|
||||
}
|
||||
|
||||
func EnqueueCaptainDocumentScheduleSyncs(ctx context.Context, wp *worker.WorkerPool, scheduledAt time.Time) (*model.BackgroundJob, error) {
|
||||
if wp == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if scheduledAt.IsZero() {
|
||||
scheduledAt = time.Now()
|
||||
}
|
||||
return wp.Enqueue(ctx, TaskTypeCaptainDocumentScheduleSyncs, captainDocumentScheduleSyncsJob{},
|
||||
worker.WithQueue("scheduled_jobs"),
|
||||
worker.WithScheduledAt(scheduledAt),
|
||||
worker.WithMaxAttempts(3),
|
||||
worker.WithIdempotencyKey(captainDocumentScheduleIdempotencyKey(scheduledAt)),
|
||||
)
|
||||
}
|
||||
|
||||
func captainDocumentScheduleIdempotencyKey(scheduledAt time.Time) string {
|
||||
bucket := scheduledAt.UTC().Truncate(captainDocumentScheduleInterval).Unix()
|
||||
return fmt.Sprintf("captain:documents_schedule_syncs:%d", bucket)
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) performDocumentSyncJob(ctx context.Context, job *model.BackgroundJob) error {
|
||||
@@ -43,3 +89,41 @@ func (s *CaptainDocumentService) performDocumentSyncJob(ctx context.Context, job
|
||||
_, err := s.SyncDocumentByAccount(ctx, payload.AccountID, payload.DocumentID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) performDocumentCrawlJob(ctx context.Context, job *model.BackgroundJob) error {
|
||||
var payload captainDocumentCrawlJob
|
||||
if err := json.Unmarshal(job.Payload, &payload); err != nil {
|
||||
return fmt.Errorf("unmarshal captain document crawl job: %w", err)
|
||||
}
|
||||
if payload.AccountID == 0 || payload.DocumentID == 0 {
|
||||
return fmt.Errorf("invalid captain document crawl job payload: %#v", payload)
|
||||
}
|
||||
_, err := s.CrawlDocumentByAccount(ctx, payload.AccountID, payload.DocumentID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) performDocumentPageCrawlParseJob(ctx context.Context, job *model.BackgroundJob) error {
|
||||
var payload captainDocumentPageCrawlParseJob
|
||||
if err := json.Unmarshal(job.Payload, &payload); err != nil {
|
||||
return fmt.Errorf("unmarshal captain document page crawl parse job: %w", err)
|
||||
}
|
||||
if payload.AccountID == 0 || payload.AssistantID == 0 || payload.PageLink == "" {
|
||||
return fmt.Errorf("invalid captain document page crawl parse job payload: %#v", payload)
|
||||
}
|
||||
_, err := s.ParseCrawledPage(ctx, payload.AccountID, payload.AssistantID, payload.PageLink)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) performDocumentScheduleSyncsJob(ctx context.Context, job *model.BackgroundJob) error {
|
||||
if len(job.Payload) > 0 {
|
||||
var payload captainDocumentScheduleSyncsJob
|
||||
if err := json.Unmarshal(job.Payload, &payload); err != nil {
|
||||
return fmt.Errorf("unmarshal captain document schedule syncs job: %w", err)
|
||||
}
|
||||
}
|
||||
if _, err := s.ScheduleDueDocumentSyncs(ctx, time.Now()); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := EnqueueCaptainDocumentScheduleSyncs(ctx, s.worker, time.Now().Add(captainDocumentScheduleInterval))
|
||||
return err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user