Enhance company recommendation caching and pipeline integration
continuous-integration/drone/push Build is passing

- Updated the company service to include a new method for scheduling the refresh of cached AI recommendations after the AI pipeline execution.
- Introduced a new interface for managing cached recommendation refreshes, improving the separation of concerns within the service layer.
- Enhanced the worker initialization to include Redis client support, allowing for better management of recommendation caching.
- Added functionality to list company IDs with existing recommendation caches, ensuring efficient updates post-pipeline runs.
- Implemented unit tests to validate the new recommendation refresh logic and ensure proper handling of various scenarios.

This update significantly improves the handling of AI recommendations by integrating caching mechanisms with the AI pipeline, enhancing overall system performance and responsiveness.
This commit is contained in:
Mazyar
2026-07-08 00:45:53 +03:30
parent 4e5296d5dd
commit 0e4fadaf29
9 changed files with 421 additions and 25 deletions
+32 -7
View File
@@ -48,23 +48,35 @@ type Service interface {
GetMinioStats(ctx context.Context) (ai_summarizer.PipelineMinioStatsResponse, error)
}
// CachedRecommendationRefresher schedules a background refresh of company recommendation caches.
type CachedRecommendationRefresher interface {
ScheduleRefreshCachedAIRecommendationsAfterPipeline()
}
// ScrapedDocumentMetadataSyncer persists MinIO scraped-document metadata onto tender records.
type ScrapedDocumentMetadataSyncer interface {
SyncScrapedDocumentsFromStorage(ctx context.Context, contractFolderID, noticePublicationID string) error
}
type service struct {
client Client
logger logger.Logger
metadataSyncer ScrapedDocumentMetadataSyncer
client Client
logger logger.Logger
metadataSyncer ScrapedDocumentMetadataSyncer
recommendationRefresher CachedRecommendationRefresher
}
// NewService creates an AI pipeline admin service.
func NewService(client Client, log logger.Logger, metadataSyncer ScrapedDocumentMetadataSyncer) Service {
func NewService(
client Client,
log logger.Logger,
metadataSyncer ScrapedDocumentMetadataSyncer,
recommendationRefresher CachedRecommendationRefresher,
) Service {
return &service{
client: client,
logger: log,
metadataSyncer: metadataSyncer,
client: client,
logger: log,
metadataSyncer: metadataSyncer,
recommendationRefresher: recommendationRefresher,
}
}
@@ -219,6 +231,8 @@ func (s *service) Sync(ctx context.Context) (*ai_summarizer.PipelineSyncResponse
return nil, fmt.Errorf("pipeline sync failed: %w", err)
}
s.scheduleRecommendationRefreshAfterPipeline()
return resp, nil
}
@@ -322,6 +336,8 @@ func (s *service) DailyRun(ctx context.Context) (*ai_summarizer.PipelineStartedR
return nil, fmt.Errorf("pipeline daily-run failed: %w", err)
}
s.scheduleRecommendationRefreshAfterPipeline()
return resp, nil
}
@@ -340,6 +356,8 @@ func (s *service) Auto(ctx context.Context) (*ai_summarizer.PipelineStartedRespo
return nil, fmt.Errorf("pipeline auto failed: %w", err)
}
s.scheduleRecommendationRefreshAfterPipeline()
return resp, nil
}
@@ -425,6 +443,13 @@ func (s *service) GetMinioStats(ctx context.Context) (ai_summarizer.PipelineMini
return resp, nil
}
func (s *service) scheduleRecommendationRefreshAfterPipeline() {
if s.recommendationRefresher == nil {
return
}
s.recommendationRefresher.ScheduleRefreshCachedAIRecommendationsAfterPipeline()
}
func validateTenderBatch(forms []TenderRefForm) ([]ai_summarizer.TenderRef, error) {
if len(forms) == 0 {
return nil, ErrEmptyTenderBatch
@@ -0,0 +1,221 @@
package company
import (
"context"
"strings"
"sync/atomic"
"time"
"golang.org/x/sync/errgroup"
"tm/pkg/ai_summarizer"
)
// AIPipelineStatusClient reports pipeline job status from the AI service.
type AIPipelineStatusClient interface {
GetPipelineLastRun(ctx context.Context) (*ai_summarizer.PipelineReportResponse, error)
}
const (
recommendationPipelineWaitInitialDelay = 5 * time.Minute
recommendationPipelineWaitMaxAttempts = 24
recommendationRefreshConcurrency = 3
)
// ScheduleRefreshCachedAIRecommendationsAfterPipeline waits for the AI daily pipeline to finish,
// then re-fetches ranked tenders for every company that already has a recommendation cache.
func (s *companyService) ScheduleRefreshCachedAIRecommendationsAfterPipeline() {
if s.aiRecommendationClient == nil {
return
}
go func() {
ctx := context.Background()
s.logger.Info("Scheduling AI recommendation cache refresh after pipeline sync", map[string]interface{}{})
if err := s.refreshCachedAIRecommendationsAfterPipeline(ctx); err != nil {
s.logger.Error("AI recommendation cache refresh after pipeline failed", map[string]interface{}{
"error": err.Error(),
})
return
}
s.logger.Info("AI recommendation cache refresh after pipeline completed", map[string]interface{}{})
}()
}
func (s *companyService) refreshCachedAIRecommendationsAfterPipeline(ctx context.Context) error {
if err := s.waitForPipelineDailyRunCompletion(ctx); err != nil {
return err
}
companyIDs, err := s.listCompanyIDsWithRecommendationCache(ctx)
if err != nil {
return err
}
if len(companyIDs) == 0 {
s.logger.Info("No companies with recommendation cache to refresh after pipeline", map[string]interface{}{})
return nil
}
s.logger.Info("Refreshing AI recommendation caches after pipeline", map[string]interface{}{
"company_count": len(companyIDs),
})
group, groupCtx := errgroup.WithContext(ctx)
group.SetLimit(recommendationRefreshConcurrency)
var refreshed atomic.Int32
var failed atomic.Int32
for _, companyID := range companyIDs {
companyID := companyID
group.Go(func() error {
if err := s.refreshAIRecommendationsCache(groupCtx, companyID); err != nil {
s.logger.Error("Failed to refresh AI recommendation cache for company", map[string]interface{}{
"company_id": companyID,
"error": err.Error(),
})
failed.Add(1)
return nil
}
refreshed.Add(1)
return nil
})
}
if err := group.Wait(); err != nil {
return err
}
s.logger.Info("AI recommendation cache refresh summary", map[string]interface{}{
"company_count": len(companyIDs),
"refreshed": refreshed.Load(),
"failed": failed.Load(),
})
return nil
}
func (s *companyService) waitForPipelineDailyRunCompletion(ctx context.Context) error {
if s.aiPipelineStatusClient == nil {
s.logger.Info("AI pipeline status client not configured; waiting before recommendation refresh", map[string]interface{}{
"delay_sec": int(recommendationPipelineWaitInitialDelay.Seconds()),
})
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(recommendationPipelineWaitInitialDelay):
}
return nil
}
delay := recommendationPipelineWaitInitialDelay
for attempt := 1; attempt <= recommendationPipelineWaitMaxAttempts; attempt++ {
if attempt == 1 {
s.logger.Info("Waiting before first AI pipeline status check", map[string]interface{}{
"delay_sec": int(delay.Seconds()),
})
} else {
s.logger.Info("Retrying AI pipeline status check", map[string]interface{}{
"attempt": attempt,
"delay_sec": int(delay.Seconds()),
})
}
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(delay):
}
if attempt > 1 {
delay *= 2
if delay > 30*time.Minute {
delay = 30 * time.Minute
}
}
report, err := s.aiPipelineStatusClient.GetPipelineLastRun(ctx)
if err != nil {
s.logger.Warn("Failed to read AI pipeline last-run status", map[string]interface{}{
"attempt": attempt,
"error": err.Error(),
})
if attempt == recommendationPipelineWaitMaxAttempts {
s.logger.Warn("Proceeding with recommendation refresh without pipeline status", map[string]interface{}{})
return nil
}
continue
}
if !isPipelineDailyRunInProgress(report) {
s.logger.Info("AI pipeline daily-run finished; refreshing recommendation caches", map[string]interface{}{
"attempt": attempt,
"status": strings.TrimSpace(report.Status),
})
return nil
}
if attempt == recommendationPipelineWaitMaxAttempts {
s.logger.Warn("AI pipeline still in progress after max wait; refreshing recommendation caches anyway", map[string]interface{}{
"attempts": attempt,
"status": strings.TrimSpace(report.Status),
})
return nil
}
}
return nil
}
func isPipelineDailyRunInProgress(report *ai_summarizer.PipelineReportResponse) bool {
if report == nil {
return true
}
status := strings.ToLower(strings.TrimSpace(report.Status))
switch status {
case "", "running", "in_progress", "started", "processing", "pending", "queued":
return true
default:
return false
}
}
func (s *companyService) listCompanyIDsWithRecommendationCache(ctx context.Context) ([]string, error) {
allIDs, err := s.repository.ListIDs(ctx)
if err != nil {
s.logger.Error("Failed to list companies for recommendation refresh", map[string]interface{}{
"error": err.Error(),
})
return nil, err
}
cached := make([]string, 0, len(allIDs))
for _, companyID := range allIDs {
if _, ok := s.getCachedAIRecommendations(ctx, companyID); ok {
cached = append(cached, companyID)
}
}
return cached, nil
}
// refreshAIRecommendationsCache updates Redis when the AI service returns a non-empty ranked list.
// Empty or failed fetches leave the existing cache untouched.
func (s *companyService) refreshAIRecommendationsCache(ctx context.Context, companyID string) error {
responses, err := s.fetchAIRecommendations(ctx, companyID)
if err != nil {
return err
}
if len(responses) == 0 {
s.logger.Info("Skipping recommendation cache update after pipeline: AI returned empty list", map[string]interface{}{
"company_id": companyID,
})
return nil
}
s.cacheAIRecommendations(ctx, companyID, responses)
s.logger.Info("AI recommendation cache refreshed after pipeline", map[string]interface{}{
"company_id": companyID,
"count": len(responses),
})
return nil
}
@@ -0,0 +1,30 @@
package company
import (
"testing"
"tm/pkg/ai_summarizer"
)
func TestIsPipelineDailyRunInProgress(t *testing.T) {
tests := []struct {
name string
report *ai_summarizer.PipelineReportResponse
want bool
}{
{name: "nil report", report: nil, want: true},
{name: "empty status", report: &ai_summarizer.PipelineReportResponse{}, want: true},
{name: "running", report: &ai_summarizer.PipelineReportResponse{Status: "running"}, want: true},
{name: "in progress", report: &ai_summarizer.PipelineReportResponse{Status: "in_progress"}, want: true},
{name: "completed", report: &ai_summarizer.PipelineReportResponse{Status: "completed"}, want: false},
{name: "failed", report: &ai_summarizer.PipelineReportResponse{Status: "failed"}, want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := isPipelineDailyRunInProgress(tt.report); got != tt.want {
t.Fatalf("isPipelineDailyRunInProgress() = %v, want %v", got, tt.want)
}
})
}
}
+44
View File
@@ -24,6 +24,7 @@ type Repository interface {
GetByRegistrationNumber(ctx context.Context, registrationNumber string) (*Company, error)
GetByTaxID(ctx context.Context, taxID string) (*Company, error)
Search(ctx context.Context, form *SearchForm, pagination *response.Pagination) (*orm.PaginatedResult[Company], error)
ListIDs(ctx context.Context) ([]string, error)
Delete(ctx context.Context, id string) error
UpdateStatus(ctx context.Context, id string, status CompanyStatus) error
UpdateVerification(ctx context.Context, id string, isVerified, isCompliant bool, complianceNotes *string) error
@@ -383,6 +384,49 @@ func (r *companyRepository) Search(ctx context.Context, form *SearchForm, pagina
return result, nil
}
// ListIDs returns every company document ID.
func (r *companyRepository) ListIDs(ctx context.Context) ([]string, error) {
const pageSize = 500
ids := make([]string, 0)
skip := 0
for {
result, err := r.ormRepo.FindAll(ctx, bson.M{}, orm.Pagination{
Limit: pageSize,
Skip: skip,
SortField: "_id",
SortOrder: 1,
SkipCount: true,
Projection: bson.M{"_id": 1},
})
if err != nil {
r.logger.Error("Failed to list company IDs", map[string]interface{}{
"error": err.Error(),
"skip": skip,
})
return nil, err
}
if len(result.Items) == 0 {
break
}
for i := range result.Items {
id := strings.TrimSpace(result.Items[i].GetID())
if id != "" {
ids = append(ids, id)
}
}
if len(result.Items) < pageSize {
break
}
skip += pageSize
}
return ids, nil
}
// UpdateStatus updates company status
func (r *companyRepository) UpdateStatus(ctx context.Context, id string, status CompanyStatus) error {
// Get company first
+6
View File
@@ -48,6 +48,9 @@ type Service interface {
// GetAIRecommendations returns ranked tender recommendations from the AI team
GetAIRecommendations(ctx context.Context, companyID string) ([]RecommendedTenderResponse, error)
GetMergedAIRecommendations(ctx context.Context, companyIDs []string) ([]RecommendedTenderResponse, error)
// ScheduleRefreshCachedAIRecommendationsAfterPipeline re-fetches cached company recommendations after tender sync.
ScheduleRefreshCachedAIRecommendationsAfterPipeline()
}
// companyService implements the Service interface
@@ -56,6 +59,7 @@ type companyService struct {
categoryService company_category.Service
fileStore filestore.FileStoreService
aiRecommendationClient AIRecommendationClient
aiPipelineStatusClient AIPipelineStatusClient
redisClient redis.Client
recommendationCacheTTL time.Duration
logger logger.Logger
@@ -67,6 +71,7 @@ func NewService(
categoryService company_category.Service,
fileStore filestore.FileStoreService,
aiRecommendationClient AIRecommendationClient,
aiPipelineStatusClient AIPipelineStatusClient,
redisClient redis.Client,
recommendationCacheTTL time.Duration,
logger logger.Logger,
@@ -76,6 +81,7 @@ func NewService(
categoryService: categoryService,
fileStore: fileStore,
aiRecommendationClient: aiRecommendationClient,
aiPipelineStatusClient: aiPipelineStatusClient,
redisClient: redisClient,
recommendationCacheTTL: recommendationCacheTTL,
logger: logger,