85 lines
3.2 KiB
Go
85 lines
3.2 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/gochat/gochat/internal/model"
|
|
"github.com/pgvector/pgvector-go"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// CaptainAssistantResponseRepo provides data access for CaptainAssistantResponse.
|
|
type CaptainAssistantResponseRepo struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) DB() *gorm.DB {
|
|
return r.db
|
|
}
|
|
|
|
func NewCaptainAssistantResponseRepo(db *gorm.DB) *CaptainAssistantResponseRepo {
|
|
return &CaptainAssistantResponseRepo{db: db}
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) Create(ctx context.Context, resp *model.CaptainAssistantResponse) error {
|
|
return r.db.WithContext(ctx).Create(resp).Error
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) GetByID(ctx context.Context, id uint) (*model.CaptainAssistantResponse, error) {
|
|
var resp model.CaptainAssistantResponse
|
|
if err := r.db.WithContext(ctx).First(&resp, id).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
return &resp, nil
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) Update(ctx context.Context, resp *model.CaptainAssistantResponse) error {
|
|
return r.db.WithContext(ctx).Save(resp).Error
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) Delete(ctx context.Context, id uint) error {
|
|
return r.db.WithContext(ctx).Delete(&model.CaptainAssistantResponse{}, id).Error
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) ListByAssistant(ctx context.Context, assistantID uint, offset, limit int) ([]model.CaptainAssistantResponse, int64, error) {
|
|
var responses []model.CaptainAssistantResponse
|
|
var count int64
|
|
db := r.db.WithContext(ctx).Model(&model.CaptainAssistantResponse{}).Where("assistant_id = ?", assistantID)
|
|
db.Count(&count)
|
|
if err := db.Offset(offset).Limit(limit).Find(&responses).Error; err != nil {
|
|
return nil, 0, err
|
|
}
|
|
return responses, count, nil
|
|
}
|
|
|
|
func (r *CaptainAssistantResponseRepo) ListByDocument(ctx context.Context, documentableID uint, documentableType string, offset, limit int) ([]model.CaptainAssistantResponse, int64, error) {
|
|
var responses []model.CaptainAssistantResponse
|
|
var count int64
|
|
db := r.db.WithContext(ctx).Model(&model.CaptainAssistantResponse{}).
|
|
Where("documentable_id = ? AND documentable_type = ?", documentableID, documentableType)
|
|
db.Count(&count)
|
|
if err := db.Offset(offset).Limit(limit).Find(&responses).Error; err != nil {
|
|
return nil, 0, err
|
|
}
|
|
return responses, count, nil
|
|
}
|
|
|
|
// SimilaritySearch performs pgvector cosine similarity search.
|
|
// Reference: Chatwoot Captain::AssistantResponsesSearchService
|
|
func (r *CaptainAssistantResponseRepo) SimilaritySearch(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error) {
|
|
var responses []model.CaptainAssistantResponse
|
|
// Cosine distance (<=>) orders by closest first.
|
|
if err := r.db.WithContext(ctx).
|
|
Where("assistant_id = ? AND status = ?", assistantID, model.ResponseStatusApproved).
|
|
Order(gorm.Expr("embedding <=> ?", embedding)).
|
|
Limit(limit).
|
|
Find(&responses).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
return responses, nil
|
|
}
|
|
|
|
// SearchByEmbedding is an alias for SimilaritySearch for backward compatibility.
|
|
func (r *CaptainAssistantResponseRepo) SearchByEmbedding(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error) {
|
|
return r.SimilaritySearch(ctx, assistantID, embedding, limit)
|
|
} |