Files
agent/backend/internal/logic/search/service.go

175 lines
5.3 KiB
Go

package search
import (
"fmt"
"sort"
"strings"
"time"
"gorm.io/gorm"
"senlinai-agent/backend/internal/models"
)
type Service struct {
db *gorm.DB
}
const (
maxSearchResults = 50
maxSearchSnippetRunes = 240
)
type rankedSearchResult struct {
result SearchResultDTO
updatedAt time.Time
}
// SearchResultDTO 是搜索接口的稳定结果,所有关联均使用公开 identity。
type SearchResultDTO struct {
Type string `json:"type"`
ID string `json:"id"`
ProjectID string `json:"projectId"`
Title string `json:"title"`
Snippet string `json:"snippet"`
}
func NewService(databases ...*gorm.DB) *Service {
var database *gorm.DB
if len(databases) > 0 {
database = databases[0]
}
return &Service{db: database}
}
func (s *Service) database() *gorm.DB {
if s.db != nil {
return s.db
}
return models.DBService
}
func (s *Service) Search(userID uint, query string) ([]SearchResultDTO, error) {
query = strings.TrimSpace(query)
if query == "" {
return []SearchResultDTO{}, nil
}
like := containsPattern(query)
operator := "LIKE"
if s.database().Dialector.Name() == "postgres" {
operator = "ILIKE"
}
var projects []models.SaProject
projectFilter := fmt.Sprintf("owner_id = ? AND (name %s ? ESCAPE '!' OR description %s ? ESCAPE '!')", operator, operator)
if err := s.database().Where(projectFilter, userID, like, like).
Order("updated_at DESC").Order("identity ASC").
Limit(maxSearchResults).Find(&projects).Error; err != nil {
return nil, err
}
projectResults := make([]rankedSearchResult, 0, len(projects))
for _, project := range projects {
projectResults = append(projectResults, rankedSearchResult{
result: SearchResultDTO{Type: "project", ID: project.Identity, ProjectID: project.Identity, Title: project.Name, Snippet: project.Description},
updatedAt: project.UpdatedAt,
})
}
var tasks []models.SaTask
taskFilter := fmt.Sprintf("(sa_projects.owner_id = ? OR sa_tasks.assignee_id = ?) AND (sa_tasks.title %s ? ESCAPE '!' OR sa_tasks.description %s ? ESCAPE '!')", operator, operator)
if err := s.database().Select("sa_tasks.*").Joins("JOIN sa_projects ON sa_projects.id = sa_tasks.project_id").
Where(taskFilter, userID, userID, like, like).
Order("sa_tasks.updated_at DESC").Order("sa_tasks.identity ASC").
Limit(maxSearchResults).
Find(&tasks).Error; err != nil {
return nil, err
}
taskResults := make([]rankedSearchResult, 0, len(tasks))
for _, task := range tasks {
taskResults = append(taskResults, rankedSearchResult{
result: SearchResultDTO{Type: "task", ID: task.Identity, ProjectID: task.ProjectIdentity, Title: task.Title, Snippet: task.Description},
updatedAt: task.UpdatedAt,
})
}
var documents []models.SaDocument
documentFilter := fmt.Sprintf(`(sa_projects.owner_id = ? OR EXISTS (
SELECT 1 FROM sa_task_shares
JOIN sa_tasks ON sa_tasks.id = sa_task_shares.task_id
WHERE sa_task_shares.object_type = 'document'
AND sa_task_shares.object_id = sa_documents.id
AND sa_tasks.project_id = sa_documents.project_id
AND sa_tasks.assignee_id = ?
)) AND (sa_documents.name %s ? ESCAPE '!' OR EXISTS (
SELECT 1 FROM sa_document_contents
WHERE sa_document_contents.document_id = sa_documents.id
AND sa_document_contents.markdown %s ? ESCAPE '!'
))`, operator, operator)
if err := s.database().Select("sa_documents.*").Joins("JOIN sa_projects ON sa_projects.id = sa_documents.project_id").
Where(documentFilter, userID, userID, like, like).
Order("sa_documents.updated_at DESC").Order("sa_documents.identity ASC").
Limit(maxSearchResults).
Find(&documents).Error; err != nil {
return nil, err
}
documentResults := make([]rankedSearchResult, 0, len(documents))
for _, document := range documents {
var content models.SaDocumentContent
_ = s.database().Where("document_id = ?", document.ID).First(&content).Error
documentResults = append(documentResults, rankedSearchResult{
result: SearchResultDTO{
Type: "document", ID: document.Identity, ProjectID: document.ProjectIdentity,
Title: document.Name, Snippet: content.Markdown,
},
updatedAt: document.UpdatedAt,
})
}
return stableFairLimit(projectResults, taskResults, documentResults), nil
}
func stableFairLimit(buckets ...[]rankedSearchResult) []SearchResultDTO {
for _, bucket := range buckets {
sort.Slice(bucket, func(left, right int) bool {
if bucket[left].updatedAt.Equal(bucket[right].updatedAt) {
return bucket[left].result.ID < bucket[right].result.ID
}
return bucket[left].updatedAt.After(bucket[right].updatedAt)
})
}
results := make([]SearchResultDTO, 0, maxSearchResults)
for position := 0; len(results) < maxSearchResults; position++ {
added := false
for _, bucket := range buckets {
if position >= len(bucket) {
continue
}
result := bucket[position].result
result.Snippet = boundedSnippet(result.Snippet)
results = append(results, result)
added = true
if len(results) == maxSearchResults {
return results
}
}
if !added {
return results
}
}
return results
}
func boundedSnippet(value string) string {
runes := []rune(value)
if len(runes) <= maxSearchSnippetRunes {
return value
}
return string(runes[:maxSearchSnippetRunes])
}
func containsPattern(query string) string {
escaped := strings.NewReplacer("!", "!!", "%", "!%", "_", "!_").Replace(query)
return "%" + escaped + "%"
}