feat(captain): queue response embedding jobs
This commit is contained in:
@@ -581,6 +581,7 @@ func Bootstrap(env string) (*App, error) {
|
||||
// Captain services (P10 M10 — Captain AI + Copilot)
|
||||
captainAssistantService := service.NewCaptainAssistantService(captainAssistantRepo, captainInboxRepo, captainDocumentRepo, captainAssistantResponseRepo, llmProvider)
|
||||
captainDocumentService := service.NewCaptainDocumentService(captainDocumentRepo, llmProvider, captainAssistantRepo)
|
||||
captainDocumentService.SetResponseRepo(captainAssistantResponseRepo)
|
||||
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)
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/gochat/gochat/internal/repository"
|
||||
"github.com/gochat/gochat/internal/worker"
|
||||
applogger "github.com/gochat/gochat/pkg/logger"
|
||||
"github.com/pgvector/pgvector-go"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
@@ -23,10 +24,13 @@ import (
|
||||
type CaptainDocumentService struct {
|
||||
documentRepo *repository.CaptainDocumentRepo
|
||||
assistantRepo *repository.CaptainAssistantRepo
|
||||
responseRepo *repository.CaptainAssistantResponseRepo
|
||||
llmProvider llm.Provider
|
||||
syncBackend CaptainDocumentSyncBackend
|
||||
crawlBackend CaptainDocumentCrawlBackend
|
||||
pageParser CaptainDocumentPageParserBackend
|
||||
faqBackend CaptainDocumentFAQBackend
|
||||
embeddings CaptainDocumentEmbeddingBackend
|
||||
worker *worker.WorkerPool
|
||||
}
|
||||
|
||||
@@ -53,6 +57,19 @@ type CaptainDocumentPageParserBackend interface {
|
||||
ParseCaptainDocumentPage(ctx context.Context, pageLink string) (*CaptainDocumentSyncResult, error)
|
||||
}
|
||||
|
||||
type CaptainDocumentFAQBackend interface {
|
||||
GenerateCaptainDocumentFAQs(ctx context.Context, doc *model.CaptainDocument) ([]CaptainDocumentFAQ, error)
|
||||
}
|
||||
|
||||
type CaptainDocumentFAQ struct {
|
||||
Question string
|
||||
Answer string
|
||||
}
|
||||
|
||||
type CaptainDocumentEmbeddingBackend interface {
|
||||
GenerateCaptainEmbedding(ctx context.Context, accountID uint, content string) (pgvector.Vector, error)
|
||||
}
|
||||
|
||||
// NewCaptainDocumentService creates a new CaptainDocumentService.
|
||||
func NewCaptainDocumentService(
|
||||
documentRepo *repository.CaptainDocumentRepo,
|
||||
@@ -81,6 +98,18 @@ func (s *CaptainDocumentService) SetPageParserBackend(pageParser CaptainDocument
|
||||
s.pageParser = pageParser
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetResponseRepo(responseRepo *repository.CaptainAssistantResponseRepo) {
|
||||
s.responseRepo = responseRepo
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetFAQBackend(faqBackend CaptainDocumentFAQBackend) {
|
||||
s.faqBackend = faqBackend
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetEmbeddingBackend(embeddingBackend CaptainDocumentEmbeddingBackend) {
|
||||
s.embeddings = embeddingBackend
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) SetWorkerPool(wp *worker.WorkerPool) {
|
||||
s.worker = wp
|
||||
RegisterCaptainDocumentJobs(wp, s)
|
||||
@@ -425,6 +454,9 @@ func (s *CaptainDocumentService) ParseCrawledPage(ctx context.Context, accountID
|
||||
} else if err := s.documentRepo.Update(ctx, doc); err != nil {
|
||||
return nil, fmt.Errorf("update crawled document: %w", err)
|
||||
}
|
||||
if err := s.enqueueDocumentResponseBuilder(ctx, accountID, doc.ID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.documentRepo.GetByAccountAndID(ctx, accountID, doc.ID)
|
||||
}
|
||||
|
||||
@@ -472,9 +504,138 @@ func (s *CaptainDocumentService) SyncDocumentByAccount(ctx context.Context, acco
|
||||
if err := s.documentRepo.Update(ctx, doc); err != nil {
|
||||
return nil, fmt.Errorf("update synced document: %w", err)
|
||||
}
|
||||
if err := s.enqueueDocumentResponseBuilder(ctx, accountID, id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return s.documentRepo.GetByAccountAndID(ctx, accountID, id)
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) enqueueDocumentResponseBuilder(ctx context.Context, accountID, id uint) error {
|
||||
if s.worker == nil || s.responseRepo == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := s.worker.Enqueue(ctx, TaskTypeCaptainDocumentResponseBuilder, captainDocumentResponseBuilderJob{AccountID: accountID, DocumentID: id},
|
||||
worker.WithQueue("low"),
|
||||
worker.WithMaxAttempts(3),
|
||||
worker.WithIdempotencyKey(fmt.Sprintf("captain:document_response_builder:%d:%d", accountID, id)),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) BuildResponsesForDocumentByAccount(ctx context.Context, accountID, id uint) ([]model.CaptainAssistantResponse, error) {
|
||||
if s.responseRepo == nil {
|
||||
return nil, fmt.Errorf("captain response repository is required")
|
||||
}
|
||||
doc, err := s.documentRepo.GetByAccountAndID(ctx, accountID, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("document not found: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(doc.Content) == "" || doc.Status != model.DocumentStatusCompleted {
|
||||
return nil, fmt.Errorf("document is not ready for response building")
|
||||
}
|
||||
if err := s.responseRepo.DB().WithContext(ctx).
|
||||
Where("account_id = ? AND documentable_id = ? AND documentable_type IN ? AND edited = ?", accountID, doc.ID, []string{"Captain::Document", "CaptainDocument"}, false).
|
||||
Delete(&model.CaptainAssistantResponse{}).Error; err != nil {
|
||||
return nil, fmt.Errorf("reset previous document responses: %w", err)
|
||||
}
|
||||
if s.faqBackend == nil {
|
||||
return nil, fmt.Errorf("faq generation disabled")
|
||||
}
|
||||
faqs, err := s.faqBackend.GenerateCaptainDocumentFAQs(ctx, doc)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate document faqs: %w", err)
|
||||
}
|
||||
|
||||
created := make([]model.CaptainAssistantResponse, 0, len(faqs))
|
||||
for _, faq := range faqs {
|
||||
question := strings.TrimSpace(faq.Question)
|
||||
answer := strings.TrimSpace(faq.Answer)
|
||||
if question == "" || answer == "" {
|
||||
continue
|
||||
}
|
||||
documentID := doc.ID
|
||||
resp := &model.CaptainAssistantResponse{
|
||||
AccountID: accountID,
|
||||
AssistantID: doc.AssistantID,
|
||||
DocumentableID: &documentID,
|
||||
DocumentableType: "Captain::Document",
|
||||
Question: question,
|
||||
Answer: answer,
|
||||
Status: model.ResponseStatusApproved,
|
||||
Edited: false,
|
||||
}
|
||||
if err := s.responseRepo.Create(ctx, resp); err != nil {
|
||||
return created, fmt.Errorf("create document response: %w", err)
|
||||
}
|
||||
created = append(created, *resp)
|
||||
if err := s.enqueueResponseEmbedding(ctx, accountID, resp.ID, fmt.Sprintf("%s: %s", question, answer)); err != nil {
|
||||
return created, err
|
||||
}
|
||||
}
|
||||
return created, nil
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) enqueueResponseEmbedding(ctx context.Context, accountID, responseID uint, content string) error {
|
||||
if s.worker == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := s.worker.Enqueue(ctx, TaskTypeCaptainLLMUpdateEmbedding, captainLLMUpdateEmbeddingJob{AccountID: accountID, ResponseID: responseID, Content: content},
|
||||
worker.WithQueue("low"),
|
||||
worker.WithMaxAttempts(3),
|
||||
worker.WithIdempotencyKey(fmt.Sprintf("captain:llm_update_embedding:response:%d:%d", accountID, responseID)),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) UpdateAssistantResponseEmbeddingByAccount(ctx context.Context, accountID, responseID uint, content string) (*model.CaptainAssistantResponse, error) {
|
||||
if s.responseRepo == nil {
|
||||
return nil, fmt.Errorf("captain response repository is required")
|
||||
}
|
||||
resp, err := s.responseRepo.GetByAccountAndID(ctx, accountID, responseID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("response not found: %w", err)
|
||||
}
|
||||
if strings.TrimSpace(content) == "" {
|
||||
content = fmt.Sprintf("%s: %s", resp.Question, resp.Answer)
|
||||
}
|
||||
embedding, err := s.generateResponseEmbedding(ctx, accountID, content)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp.Embedding = embedding
|
||||
if s.responseRepo.DB().Dialector != nil && s.responseRepo.DB().Dialector.Name() == "sqlite" {
|
||||
if err := s.responseRepo.DB().WithContext(ctx).Omit("Embedding").Save(resp).Error; err != nil {
|
||||
return nil, fmt.Errorf("update response embedding: %w", err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
if err := s.responseRepo.Update(ctx, resp); err != nil {
|
||||
return nil, fmt.Errorf("update response embedding: %w", err)
|
||||
}
|
||||
return s.responseRepo.GetByAccountAndID(ctx, accountID, responseID)
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) generateResponseEmbedding(ctx context.Context, accountID uint, content string) (pgvector.Vector, error) {
|
||||
if s.embeddings != nil {
|
||||
return s.embeddings.GenerateCaptainEmbedding(ctx, accountID, content)
|
||||
}
|
||||
if s.llmProvider == nil {
|
||||
return pgvector.Vector{}, fmt.Errorf("embedding generation disabled")
|
||||
}
|
||||
result, err := s.llmProvider.CreateEmbedding(ctx, llm.EmbeddingRequest{Model: "text-embedding-3-small", Input: []string{content}})
|
||||
if err != nil {
|
||||
return pgvector.Vector{}, fmt.Errorf("generate response embedding: %w", err)
|
||||
}
|
||||
if len(result.Data) == 0 {
|
||||
return pgvector.Vector{}, fmt.Errorf("no embedding returned")
|
||||
}
|
||||
values := make([]float32, len(result.Data[0].Embedding))
|
||||
for i, value := range result.Data[0].Embedding {
|
||||
values[i] = float32(value)
|
||||
}
|
||||
return pgvector.NewVector(values), nil
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) markDocumentSyncStarted(ctx context.Context, doc *model.CaptainDocument) error {
|
||||
now := time.Now().Unix()
|
||||
doc.LastSyncAttemptedAt = &now
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"github.com/gochat/gochat/internal/model"
|
||||
"github.com/gochat/gochat/internal/repository"
|
||||
"github.com/gochat/gochat/internal/worker"
|
||||
"github.com/pgvector/pgvector-go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
@@ -21,14 +22,17 @@ func setupCaptainDocumentServiceTest(t *testing.T) (*gorm.DB, *CaptainDocumentSe
|
||||
dbName := fmt.Sprintf("file:%s?mode=memory&cache=private", t.Name())
|
||||
db, err := gorm.Open(sqlite.Open(dbName), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&model.Account{}, &model.CaptainAssistant{}, &model.CaptainDocument{}, &model.BackgroundJob{}))
|
||||
require.NoError(t, db.AutoMigrate(&model.Account{}, &model.CaptainAssistant{}, &model.CaptainDocument{}, &model.CaptainAssistantResponse{}, &model.BackgroundJob{}))
|
||||
t.Cleanup(func() {
|
||||
sqlDB, _ := db.DB()
|
||||
sqlDB.Close()
|
||||
})
|
||||
documentRepo := repository.NewCaptainDocumentRepo(db)
|
||||
assistantRepo := repository.NewCaptainAssistantRepo(db)
|
||||
return db, NewCaptainDocumentService(documentRepo, nil, assistantRepo)
|
||||
responseRepo := repository.NewCaptainAssistantResponseRepo(db)
|
||||
svc := NewCaptainDocumentService(documentRepo, nil, assistantRepo)
|
||||
svc.SetResponseRepo(responseRepo)
|
||||
return db, svc
|
||||
}
|
||||
|
||||
func seedCaptainDocumentSyncFixture(t *testing.T, db *gorm.DB) (*model.Account, *model.Account, *model.CaptainDocument) {
|
||||
@@ -266,6 +270,171 @@ func TestEnqueueCaptainDocumentScheduleSyncsUsesDailyIdempotency(t *testing.T) {
|
||||
assert.Equal(t, int64(1), count)
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_ResponseBuilderCreatesResponsesAndEmbeddingJobs(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
doc.Content = "How refunds work. How billing works."
|
||||
doc.Status = model.DocumentStatusCompleted
|
||||
require.NoError(t, db.Save(doc).Error)
|
||||
uneditedDocID := doc.ID
|
||||
unedited := &model.CaptainAssistantResponse{
|
||||
AccountID: account.ID,
|
||||
AssistantID: doc.AssistantID,
|
||||
DocumentableID: &uneditedDocID,
|
||||
DocumentableType: "Captain::Document",
|
||||
Question: "old",
|
||||
Answer: "old answer",
|
||||
Status: model.ResponseStatusApproved,
|
||||
Edited: false,
|
||||
}
|
||||
edited := &model.CaptainAssistantResponse{
|
||||
AccountID: account.ID,
|
||||
AssistantID: doc.AssistantID,
|
||||
DocumentableID: &uneditedDocID,
|
||||
DocumentableType: "Captain::Document",
|
||||
Question: "keep",
|
||||
Answer: "edited answer",
|
||||
Status: model.ResponseStatusApproved,
|
||||
Edited: true,
|
||||
}
|
||||
require.NoError(t, db.Create(unedited).Error)
|
||||
require.NoError(t, db.Create(edited).Error)
|
||||
svc.SetFAQBackend(&captainDocumentFakeFAQBackend{faqs: []CaptainDocumentFAQ{
|
||||
{Question: "How do refunds work?", Answer: "Refunds take five days."},
|
||||
{Question: "How does billing work?", Answer: "Billing is monthly."},
|
||||
}})
|
||||
now := time.Date(2026, 6, 6, 2, 0, 0, 0, time.UTC)
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithNow(func() time.Time { return now }))
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
_, err := wp.Enqueue(context.Background(), TaskTypeCaptainDocumentResponseBuilder, captainDocumentResponseBuilderJob{AccountID: account.ID, DocumentID: doc.ID}, worker.WithQueue("low"), worker.WithMaxAttempts(3))
|
||||
require.NoError(t, err)
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var deleted model.CaptainAssistantResponse
|
||||
assert.Error(t, db.First(&deleted, unedited.ID).Error)
|
||||
var kept model.CaptainAssistantResponse
|
||||
require.NoError(t, db.First(&kept, edited.ID).Error)
|
||||
assert.True(t, kept.Edited)
|
||||
|
||||
var responses []model.CaptainAssistantResponse
|
||||
require.NoError(t, db.Where("account_id = ? AND documentable_id = ? AND edited = ?", account.ID, doc.ID, false).Order("id ASC").Find(&responses).Error)
|
||||
require.Len(t, responses, 2)
|
||||
assert.Equal(t, "Captain::Document", responses[0].DocumentableType)
|
||||
assert.Equal(t, model.ResponseStatusApproved, responses[0].Status)
|
||||
|
||||
var embeddingJobs []model.BackgroundJob
|
||||
require.NoError(t, db.Where("job_type = ? AND queue = ?", TaskTypeCaptainLLMUpdateEmbedding, "low").Find(&embeddingJobs).Error)
|
||||
require.Len(t, embeddingJobs, 2)
|
||||
assert.Contains(t, string(embeddingJobs[0].Payload), fmt.Sprintf("\"response_id\":%d", responses[0].ID))
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_EmbeddingJobUpdatesResponse(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
documentID := doc.ID
|
||||
resp := &model.CaptainAssistantResponse{
|
||||
AccountID: account.ID,
|
||||
AssistantID: doc.AssistantID,
|
||||
DocumentableID: &documentID,
|
||||
DocumentableType: "Captain::Document",
|
||||
Question: "What is Captain?",
|
||||
Answer: "Captain answers customers.",
|
||||
Status: model.ResponseStatusApproved,
|
||||
}
|
||||
require.NoError(t, db.Create(resp).Error)
|
||||
embeddings := &captainDocumentFakeEmbeddingBackend{embedding: pgvector.NewVector([]float32{0.1, 0.2, 0.3})}
|
||||
svc.SetEmbeddingBackend(embeddings)
|
||||
wp := worker.NewWorkerPool(db)
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
_, err := wp.Enqueue(context.Background(), TaskTypeCaptainLLMUpdateEmbedding, captainLLMUpdateEmbeddingJob{AccountID: account.ID, ResponseID: resp.ID}, worker.WithQueue("low"), worker.WithMaxAttempts(3))
|
||||
require.NoError(t, err)
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
assert.Equal(t, fmt.Sprintf("%s: %s", resp.Question, resp.Answer), embeddings.content)
|
||||
var job model.BackgroundJob
|
||||
require.NoError(t, db.Where("job_type = ?", TaskTypeCaptainLLMUpdateEmbedding).First(&job).Error)
|
||||
assert.Equal(t, model.BackgroundJobStatusCompleted, job.Status)
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_EmbeddingJobRetriesWhenProviderDisabled(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
documentID := doc.ID
|
||||
resp := &model.CaptainAssistantResponse{
|
||||
AccountID: account.ID,
|
||||
AssistantID: doc.AssistantID,
|
||||
DocumentableID: &documentID,
|
||||
DocumentableType: "Captain::Document",
|
||||
Question: "What is Captain?",
|
||||
Answer: "Captain answers customers.",
|
||||
Status: model.ResponseStatusApproved,
|
||||
}
|
||||
require.NoError(t, db.Create(resp).Error)
|
||||
now := time.Date(2026, 6, 6, 2, 15, 0, 0, time.UTC)
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithNow(func() time.Time { return now }), worker.WithBackoff(func(attempt int) time.Duration { return time.Minute }))
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
_, err := wp.Enqueue(context.Background(), TaskTypeCaptainLLMUpdateEmbedding, captainLLMUpdateEmbeddingJob{AccountID: account.ID, ResponseID: resp.ID}, worker.WithQueue("low"), worker.WithMaxAttempts(3))
|
||||
require.NoError(t, err)
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.Error(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var job model.BackgroundJob
|
||||
require.NoError(t, db.Where("job_type = ?", TaskTypeCaptainLLMUpdateEmbedding).First(&job).Error)
|
||||
assert.Equal(t, model.BackgroundJobStatusRetrying, job.Status)
|
||||
assert.Contains(t, job.LastError, "embedding generation disabled")
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_ResponseBuilderRetriesWhenFAQDisabled(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
account, _, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
doc.Content = "ready content"
|
||||
doc.Status = model.DocumentStatusCompleted
|
||||
require.NoError(t, db.Save(doc).Error)
|
||||
now := time.Date(2026, 6, 6, 2, 30, 0, 0, time.UTC)
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithNow(func() time.Time { return now }), worker.WithBackoff(func(attempt int) time.Duration { return time.Minute }))
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
_, err := wp.Enqueue(context.Background(), TaskTypeCaptainDocumentResponseBuilder, captainDocumentResponseBuilderJob{AccountID: account.ID, DocumentID: doc.ID}, worker.WithQueue("low"), worker.WithMaxAttempts(3))
|
||||
require.NoError(t, err)
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.Error(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var job model.BackgroundJob
|
||||
require.NoError(t, db.Where("job_type = ?", TaskTypeCaptainDocumentResponseBuilder).First(&job).Error)
|
||||
assert.Equal(t, model.BackgroundJobStatusRetrying, job.Status)
|
||||
assert.Contains(t, job.LastError, "faq generation disabled")
|
||||
}
|
||||
|
||||
func TestCaptainDocumentService_ResponseBuilderScopesAccount(t *testing.T) {
|
||||
db, svc := setupCaptainDocumentServiceTest(t)
|
||||
_, otherAccount, doc := seedCaptainDocumentSyncFixture(t, db)
|
||||
doc.Content = "ready content"
|
||||
doc.Status = model.DocumentStatusCompleted
|
||||
require.NoError(t, db.Save(doc).Error)
|
||||
svc.SetFAQBackend(&captainDocumentFakeFAQBackend{faqs: []CaptainDocumentFAQ{{Question: "Q", Answer: "A"}}})
|
||||
wp := worker.NewWorkerPoolWithOptions(db, worker.WithBackoff(func(attempt int) time.Duration { return time.Minute }))
|
||||
svc.SetWorkerPool(wp)
|
||||
|
||||
_, err := wp.Enqueue(context.Background(), TaskTypeCaptainDocumentResponseBuilder, captainDocumentResponseBuilderJob{AccountID: otherAccount.ID, DocumentID: doc.ID}, worker.WithQueue("low"), worker.WithMaxAttempts(3))
|
||||
require.NoError(t, err)
|
||||
processed, err := wp.ProcessOne(context.Background())
|
||||
require.Error(t, err)
|
||||
assert.True(t, processed)
|
||||
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&model.CaptainAssistantResponse{}).Where("documentable_id = ?", doc.ID).Count(&count).Error)
|
||||
assert.Equal(t, int64(0), count)
|
||||
}
|
||||
|
||||
type captainDocumentFakeSyncBackend struct {
|
||||
result *CaptainDocumentSyncResult
|
||||
err error
|
||||
@@ -305,3 +474,29 @@ func (b *captainDocumentFakePageParserBackend) ParseCaptainDocumentPage(ctx cont
|
||||
}
|
||||
return b.pages[pageLink], nil
|
||||
}
|
||||
|
||||
type captainDocumentFakeFAQBackend struct {
|
||||
faqs []CaptainDocumentFAQ
|
||||
err error
|
||||
}
|
||||
|
||||
func (b *captainDocumentFakeFAQBackend) GenerateCaptainDocumentFAQs(ctx context.Context, doc *model.CaptainDocument) ([]CaptainDocumentFAQ, error) {
|
||||
if b.err != nil {
|
||||
return nil, b.err
|
||||
}
|
||||
return b.faqs, nil
|
||||
}
|
||||
|
||||
type captainDocumentFakeEmbeddingBackend struct {
|
||||
embedding pgvector.Vector
|
||||
err error
|
||||
content string
|
||||
}
|
||||
|
||||
func (b *captainDocumentFakeEmbeddingBackend) GenerateCaptainEmbedding(ctx context.Context, accountID uint, content string) (pgvector.Vector, error) {
|
||||
b.content = content
|
||||
if b.err != nil {
|
||||
return pgvector.Vector{}, b.err
|
||||
}
|
||||
return b.embedding, nil
|
||||
}
|
||||
|
||||
@@ -12,10 +12,12 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
TaskTypeCaptainDocumentSync = "captain:document_sync"
|
||||
TaskTypeCaptainDocumentCrawl = "captain:document_crawl"
|
||||
TaskTypeCaptainDocumentPageCrawlParse = "captain:document_page_crawl_parse"
|
||||
TaskTypeCaptainDocumentScheduleSyncs = "captain:documents_schedule_syncs"
|
||||
TaskTypeCaptainDocumentSync = "captain:document_sync"
|
||||
TaskTypeCaptainDocumentCrawl = "captain:document_crawl"
|
||||
TaskTypeCaptainDocumentPageCrawlParse = "captain:document_page_crawl_parse"
|
||||
TaskTypeCaptainDocumentScheduleSyncs = "captain:documents_schedule_syncs"
|
||||
TaskTypeCaptainDocumentResponseBuilder = "captain:document_response_builder"
|
||||
TaskTypeCaptainLLMUpdateEmbedding = "captain:llm_update_embedding"
|
||||
)
|
||||
|
||||
const captainDocumentScheduleInterval = 24 * time.Hour
|
||||
@@ -40,6 +42,17 @@ type captainDocumentScheduleSyncsJob struct {
|
||||
PlanName string `json:"plan_name,omitempty"`
|
||||
}
|
||||
|
||||
type captainDocumentResponseBuilderJob struct {
|
||||
AccountID uint `json:"account_id"`
|
||||
DocumentID uint `json:"document_id"`
|
||||
}
|
||||
|
||||
type captainLLMUpdateEmbeddingJob struct {
|
||||
AccountID uint `json:"account_id"`
|
||||
ResponseID uint `json:"response_id"`
|
||||
Content string `json:"content,omitempty"`
|
||||
}
|
||||
|
||||
var captainDocumentRegistrations sync.Map
|
||||
|
||||
// RegisterCaptainDocumentJobs wires Captain::Documents::PerformSyncJob into
|
||||
@@ -56,6 +69,8 @@ func RegisterCaptainDocumentJobs(wp *worker.WorkerPool, svc *CaptainDocumentServ
|
||||
wp.Register(TaskTypeCaptainDocumentCrawl, svc.performDocumentCrawlJob)
|
||||
wp.Register(TaskTypeCaptainDocumentPageCrawlParse, svc.performDocumentPageCrawlParseJob)
|
||||
wp.Register(TaskTypeCaptainDocumentScheduleSyncs, svc.performDocumentScheduleSyncsJob)
|
||||
wp.Register(TaskTypeCaptainDocumentResponseBuilder, svc.performDocumentResponseBuilderJob)
|
||||
wp.Register(TaskTypeCaptainLLMUpdateEmbedding, svc.performCaptainLLMUpdateEmbeddingJob)
|
||||
}
|
||||
|
||||
func EnqueueCaptainDocumentScheduleSyncs(ctx context.Context, wp *worker.WorkerPool, scheduledAt time.Time) (*model.BackgroundJob, error) {
|
||||
@@ -127,3 +142,27 @@ func (s *CaptainDocumentService) performDocumentScheduleSyncsJob(ctx context.Con
|
||||
_, err := EnqueueCaptainDocumentScheduleSyncs(ctx, s.worker, time.Now().Add(captainDocumentScheduleInterval))
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) performDocumentResponseBuilderJob(ctx context.Context, job *model.BackgroundJob) error {
|
||||
var payload captainDocumentResponseBuilderJob
|
||||
if err := json.Unmarshal(job.Payload, &payload); err != nil {
|
||||
return fmt.Errorf("unmarshal captain document response builder job: %w", err)
|
||||
}
|
||||
if payload.AccountID == 0 || payload.DocumentID == 0 {
|
||||
return fmt.Errorf("invalid captain document response builder job payload: %#v", payload)
|
||||
}
|
||||
_, err := s.BuildResponsesForDocumentByAccount(ctx, payload.AccountID, payload.DocumentID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *CaptainDocumentService) performCaptainLLMUpdateEmbeddingJob(ctx context.Context, job *model.BackgroundJob) error {
|
||||
var payload captainLLMUpdateEmbeddingJob
|
||||
if err := json.Unmarshal(job.Payload, &payload); err != nil {
|
||||
return fmt.Errorf("unmarshal captain llm update embedding job: %w", err)
|
||||
}
|
||||
if payload.AccountID == 0 || payload.ResponseID == 0 {
|
||||
return fmt.Errorf("invalid captain llm update embedding job payload: %#v", payload)
|
||||
}
|
||||
_, err := s.UpdateAssistantResponseEmbeddingByAccount(ctx, payload.AccountID, payload.ResponseID, payload.Content)
|
||||
return err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user