WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
Anydoc / anydoc parser engine (push) Canceled after 0s
App / Format, vet, test, and build (push) Canceled after 0s
CLI / test (macos-latest / go 1.26) (push) Canceled after 0s
CLI / test (ubuntu-latest / go 1.26) (push) Canceled after 0s
CLI / test (windows-latest / go 1.26) (push) Canceled after 0s
Build and Push Docker Image / build-ui (push) Canceled after 0s
Build and Push Docker Image / build-docreader (push) Canceled after 0s
Build and Push Docker Image / build-sandbox (push) Canceled after 0s
Build and Push Docker Image / build-app (amd64, linux/amd64, ubuntu-latest) (push) Canceled after 0s
Build and Push Docker Image / build-app (arm64, linux/arm64, ubuntu-24.04-arm) (push) Canceled after 0s
DocReader / Python tests and Go client tests (push) Canceled after 0s
Build & Publish (dsh-weknora) / Typecheck, test, pack (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (pinned) (push) Canceled after 0s
Build & Publish (dsh-weknora) / End-to-end inside dsh (latest, advisory) (push) Canceled after 0s
Frontend / Test, type-check, and build (push) Canceled after 0s
Go Lint Cache / warm lint cache (.) (push) Canceled after 0s
Go Lint Cache / warm lint cache (cli) (push) Canceled after 0s
Go Lint Cache / warm lint cache (client) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.10) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.11) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.12) (push) Canceled after 0s
Build & Publish (mcp-server) / Test (Python 3.13) (push) Canceled after 0s
Build and Push Docker Image / merge (push) Canceled after 0s
Build & Publish (dsh-weknora) / Publish to npm (push) Canceled after 0s
Build & Publish (mcp-server) / Build sdist + wheel (push) Canceled after 0s
Build & Publish (mcp-server) / Publish to PyPI (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/stretchr/testify/require"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAgentDeletionUnbindsFusionOnlyInSameTenant(t *testing.T) {
|
||||
_, db, scope, _ := setupDigitalHumanTest(t)
|
||||
rows := []types.CustomAgent{
|
||||
{ID: "aux", TenantID: scope.TenantID, Name: "aux"},
|
||||
{ID: "primary", TenantID: scope.TenantID, Name: "primary", Config: types.CustomAgentConfig{FusionAgentIDs: []string{"aux", "other"}, SystemPrompt: "keep"}},
|
||||
{ID: "outside", TenantID: scope.TenantID + 1, Name: "outside", Config: types.CustomAgentConfig{FusionAgentIDs: []string{"aux"}}},
|
||||
}
|
||||
for _, row := range rows {
|
||||
require.NoError(t, db.Create(&row).Error)
|
||||
}
|
||||
repo := &customAgentRepository{db: db}
|
||||
_, err := repo.DeleteAgentWithDigitalHumanBindings(context.Background(), "aux", scope.TenantID)
|
||||
require.NoError(t, err)
|
||||
var primary, outside types.CustomAgent
|
||||
require.NoError(t, db.First(&primary, "id = ?", "primary").Error)
|
||||
require.NoError(t, db.First(&outside, "id = ?", "outside").Error)
|
||||
require.Equal(t, []string{"other"}, primary.Config.FusionAgentIDs)
|
||||
require.Equal(t, "keep", primary.Config.SystemPrompt)
|
||||
require.Equal(t, []string{"aux"}, outside.Config.FusionAgentIDs)
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// softDeleteAgentIMChannels participates in the caller's agent deletion
|
||||
// transaction. Runtime invalidation must happen only after that transaction
|
||||
// commits. The IM package depends on repositories, so use its durable table
|
||||
// contract here rather than importing its runtime service.
|
||||
func softDeleteAgentIMChannels(tx *gorm.DB, agentID string, tenantID uint64) ([]string, error) {
|
||||
var ids []string
|
||||
query := tx.Table("im_channels").Where("agent_id = ? AND tenant_id = ? AND deleted_at IS NULL", agentID, tenantID)
|
||||
if err := query.Pluck("id", &ids).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return ids, nil
|
||||
}
|
||||
now := time.Now()
|
||||
if err := tx.Table("im_channels").Where("id IN ? AND tenant_id = ?", ids, tenantID).
|
||||
Updates(map[string]any{"enabled": false, "deleted_at": now, "updated_at": now}).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ids, nil
|
||||
}
|
||||
@@ -0,0 +1,241 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrAgentShareNotFound = errors.New("agent share not found")
|
||||
ErrAgentShareAlreadyExists = errors.New("agent already shared to this organization")
|
||||
)
|
||||
|
||||
// agentShareRepository implements AgentShareRepository interface
|
||||
type agentShareRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewAgentShareRepository creates a new agent share repository
|
||||
func NewAgentShareRepository(db *gorm.DB) interfaces.AgentShareRepository {
|
||||
return &agentShareRepository{db: db}
|
||||
}
|
||||
|
||||
// Create creates a new agent share record
|
||||
func (r *agentShareRepository) Create(ctx context.Context, share *types.AgentShare) error {
|
||||
var count int64
|
||||
r.db.WithContext(ctx).Model(&types.AgentShare{}).
|
||||
Where("agent_id = ? AND source_tenant_id = ? AND organization_id = ? AND deleted_at IS NULL",
|
||||
share.AgentID, share.SourceTenantID, share.OrganizationID).
|
||||
Count(&count)
|
||||
if count > 0 {
|
||||
return ErrAgentShareAlreadyExists
|
||||
}
|
||||
return r.db.WithContext(ctx).Create(share).Error
|
||||
}
|
||||
|
||||
// GetByID gets a share record by ID
|
||||
func (r *agentShareRepository) GetByID(ctx context.Context, id string) (*types.AgentShare, error) {
|
||||
var share types.AgentShare
|
||||
err := r.db.WithContext(ctx).Where("id = ?", id).First(&share).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrAgentShareNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &share, nil
|
||||
}
|
||||
|
||||
// GetByAgentAndOrg gets a share record by agent ID and organization ID
|
||||
func (r *agentShareRepository) GetByAgentAndOrg(ctx context.Context, agentID string, orgID string) (*types.AgentShare, error) {
|
||||
var share types.AgentShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("agent_id = ? AND organization_id = ?", agentID, orgID).
|
||||
First(&share).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrAgentShareNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &share, nil
|
||||
}
|
||||
|
||||
// Update updates a share record
|
||||
func (r *agentShareRepository) Update(ctx context.Context, share *types.AgentShare) error {
|
||||
return r.db.WithContext(ctx).Model(&types.AgentShare{}).
|
||||
Where("id = ?", share.ID).Updates(share).Error
|
||||
}
|
||||
|
||||
// Delete soft deletes a share record
|
||||
func (r *agentShareRepository) Delete(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).Where("id = ?", id).Delete(&types.AgentShare{}).Error
|
||||
}
|
||||
|
||||
// DeleteByAgentIDAndSourceTenant soft deletes all share records for an agent (id, tenant_id)
|
||||
func (r *agentShareRepository) DeleteByAgentIDAndSourceTenant(ctx context.Context, agentID string, sourceTenantID uint64) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("agent_id = ? AND source_tenant_id = ?", agentID, sourceTenantID).
|
||||
Delete(&types.AgentShare{}).Error
|
||||
}
|
||||
|
||||
// DeleteByOrganizationID soft deletes all share records for an organization
|
||||
func (r *agentShareRepository) DeleteByOrganizationID(ctx context.Context, orgID string) error {
|
||||
return r.db.WithContext(ctx).Where("organization_id = ?", orgID).Delete(&types.AgentShare{}).Error
|
||||
}
|
||||
|
||||
// ListByAgent lists all share records for an agent
|
||||
func (r *agentShareRepository) ListByAgent(ctx context.Context, agentID string) ([]*types.AgentShare, error) {
|
||||
var shares []*types.AgentShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Preload("Organization").
|
||||
Where("agent_id = ?", agentID).
|
||||
Order("created_at DESC").
|
||||
Find(&shares).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// ListByOrganization lists all share records for an organization (excluding deleted agents)
|
||||
func (r *agentShareRepository) ListByOrganization(ctx context.Context, orgID string) ([]*types.AgentShare, error) {
|
||||
var shares []*types.AgentShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN custom_agents ON custom_agents.id = agent_shares.agent_id AND custom_agents.tenant_id = agent_shares.source_tenant_id AND custom_agents.deleted_at IS NULL").
|
||||
Preload("Agent").
|
||||
Preload("Organization").
|
||||
Where("agent_shares.organization_id = ? AND agent_shares.deleted_at IS NULL", orgID).
|
||||
Order("agent_shares.created_at DESC").
|
||||
Find(&shares).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// ListByOrganizations lists all share records for the given organizations (batch).
|
||||
func (r *agentShareRepository) ListByOrganizations(ctx context.Context, orgIDs []string) ([]*types.AgentShare, error) {
|
||||
if len(orgIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var shares []*types.AgentShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN custom_agents ON custom_agents.id = agent_shares.agent_id AND custom_agents.tenant_id = agent_shares.source_tenant_id AND custom_agents.deleted_at IS NULL").
|
||||
Preload("Agent").
|
||||
Preload("Organization").
|
||||
Where("agent_shares.organization_id IN ? AND agent_shares.deleted_at IS NULL", orgIDs).
|
||||
Order("agent_shares.created_at DESC").
|
||||
Find(&shares).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// CountByOrganizations returns share counts per organization (only orgs in orgIDs). Excludes deleted agents.
|
||||
func (r *agentShareRepository) CountByOrganizations(ctx context.Context, orgIDs []string) (map[string]int64, error) {
|
||||
if len(orgIDs) == 0 {
|
||||
return make(map[string]int64), nil
|
||||
}
|
||||
type row struct {
|
||||
OrgID string `gorm:"column:organization_id"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
}
|
||||
var rows []row
|
||||
err := r.db.WithContext(ctx).Model(&types.AgentShare{}).
|
||||
Joins("JOIN custom_agents ON custom_agents.id = agent_shares.agent_id AND custom_agents.tenant_id = agent_shares.source_tenant_id AND custom_agents.deleted_at IS NULL").
|
||||
Select("agent_shares.organization_id as organization_id, COUNT(*) as count").
|
||||
Where("agent_shares.organization_id IN ? AND agent_shares.deleted_at IS NULL", orgIDs).
|
||||
Group("agent_shares.organization_id").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]int64)
|
||||
for _, o := range orgIDs {
|
||||
out[o] = 0
|
||||
}
|
||||
for _, r := range rows {
|
||||
out[r.OrgID] = r.Count
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ListSharedAgentsForTenant lists all agents shared to organizations that the
|
||||
// caller's tenant participates in. Plan 3 of #1303 keys this on tenant rather
|
||||
// than user.
|
||||
func (r *agentShareRepository) ListSharedAgentsForTenant(ctx context.Context, tenantID uint64) ([]*types.AgentShare, error) {
|
||||
var shares []*types.AgentShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN custom_agents ON custom_agents.id = agent_shares.agent_id AND custom_agents.tenant_id = agent_shares.source_tenant_id AND custom_agents.deleted_at IS NULL").
|
||||
Preload("Agent").
|
||||
Preload("Organization").
|
||||
Joins("JOIN organization_tenant_members otm ON otm.organization_id = agent_shares.organization_id").
|
||||
Joins("JOIN organizations ON organizations.id = agent_shares.organization_id AND organizations.deleted_at IS NULL").
|
||||
Where("otm.tenant_id = ?", tenantID).
|
||||
Where("agent_shares.deleted_at IS NULL").
|
||||
Order("agent_shares.created_at DESC").
|
||||
Find(&shares).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// GetShareByAgentIDForTenant returns one share for the given agentID that the
|
||||
// tenant can reach (tenant participates in some org with the share), excluding
|
||||
// source_tenant_id == excludeTenantID. Single query.
|
||||
func (r *agentShareRepository) GetShareByAgentIDForTenant(ctx context.Context, tenantID uint64, agentID string, excludeTenantID uint64) (*types.AgentShare, error) {
|
||||
var share types.AgentShare
|
||||
tx := r.db.WithContext(ctx).
|
||||
Joins("JOIN organization_tenant_members otm ON otm.organization_id = agent_shares.organization_id").
|
||||
Where("agent_shares.agent_id = ?", agentID).
|
||||
Where("otm.tenant_id = ?", tenantID).
|
||||
Where("agent_shares.source_tenant_id != ?", excludeTenantID).
|
||||
Where("agent_shares.deleted_at IS NULL").
|
||||
Order("agent_shares.id").
|
||||
Limit(1).
|
||||
Find(&share)
|
||||
if tx.Error != nil {
|
||||
return nil, tx.Error
|
||||
}
|
||||
if tx.RowsAffected == 0 {
|
||||
return nil, ErrAgentShareNotFound
|
||||
}
|
||||
return &share, nil
|
||||
}
|
||||
|
||||
// GetShareByAgentIDAndSourceForTenant validates an exact source selector
|
||||
// against organization membership. This avoids loading every shared agent and
|
||||
// makes same-ID builtins from multiple workspaces deterministic.
|
||||
func (r *agentShareRepository) GetShareByAgentIDAndSourceForTenant(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
agentID string,
|
||||
sourceTenantID uint64,
|
||||
) (*types.AgentShare, error) {
|
||||
var share types.AgentShare
|
||||
tx := r.db.WithContext(ctx).
|
||||
Joins("JOIN organization_tenant_members otm ON otm.organization_id = agent_shares.organization_id").
|
||||
Joins("JOIN organizations ON organizations.id = agent_shares.organization_id AND organizations.deleted_at IS NULL").
|
||||
Joins("JOIN custom_agents ON custom_agents.id = agent_shares.agent_id AND custom_agents.tenant_id = agent_shares.source_tenant_id AND custom_agents.deleted_at IS NULL").
|
||||
Where("agent_shares.agent_id = ?", agentID).
|
||||
Where("agent_shares.source_tenant_id = ?", sourceTenantID).
|
||||
Where("otm.tenant_id = ?", tenantID).
|
||||
Where("agent_shares.deleted_at IS NULL").
|
||||
Order("agent_shares.id").
|
||||
Limit(1).
|
||||
Find(&share)
|
||||
if tx.Error != nil {
|
||||
return nil, tx.Error
|
||||
}
|
||||
if tx.RowsAffected == 0 {
|
||||
return nil, ErrAgentShareNotFound
|
||||
}
|
||||
return &share, nil
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestGetShareByAgentIDAndSourceForTenantDisambiguatesSource(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:agent-share-source?mode=memory&cache=shared"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
statements := []string{
|
||||
`CREATE TABLE agent_shares (
|
||||
id TEXT PRIMARY KEY, agent_id TEXT, organization_id TEXT,
|
||||
shared_by_user_id TEXT, source_tenant_id INTEGER, permission TEXT,
|
||||
created_at DATETIME, updated_at DATETIME, deleted_at DATETIME
|
||||
)`,
|
||||
`CREATE TABLE organization_tenant_members (organization_id TEXT, tenant_id INTEGER)`,
|
||||
`CREATE TABLE organizations (id TEXT PRIMARY KEY, deleted_at DATETIME)`,
|
||||
`CREATE TABLE custom_agents (id TEXT, tenant_id INTEGER, deleted_at DATETIME)`,
|
||||
`INSERT INTO organizations(id) VALUES ('org')`,
|
||||
`INSERT INTO organization_tenant_members(organization_id, tenant_id) VALUES ('org', 7)`,
|
||||
`INSERT INTO custom_agents(id, tenant_id) VALUES ('builtin-smart-reasoning', 42), ('builtin-smart-reasoning', 84)`,
|
||||
`INSERT INTO agent_shares(id, agent_id, organization_id, source_tenant_id) VALUES
|
||||
('share-42', 'builtin-smart-reasoning', 'org', 42),
|
||||
('share-84', 'builtin-smart-reasoning', 'org', 84)`,
|
||||
}
|
||||
for _, statement := range statements {
|
||||
require.NoError(t, db.Exec(statement).Error)
|
||||
}
|
||||
|
||||
repo := &agentShareRepository{db: db}
|
||||
share, err := repo.GetShareByAgentIDAndSourceForTenant(
|
||||
context.Background(), 7, "builtin-smart-reasoning", 84,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "share-84", share.ID)
|
||||
require.Equal(t, uint64(84), share.SourceTenantID)
|
||||
|
||||
_, err = repo.GetShareByAgentIDAndSourceForTenant(
|
||||
context.Background(), 8, "builtin-smart-reasoning", 84,
|
||||
)
|
||||
require.ErrorIs(t, err, ErrAgentShareNotFound)
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// auditLogRepository persists audit events to the audit_logs table.
|
||||
// Rows are append-only — there is no Update, no soft-delete, no
|
||||
// Modify hook. The cursor pagination uses the monotonic id column,
|
||||
// not created_at, so duplicate timestamps don't break iteration.
|
||||
type auditLogRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewAuditLogRepository constructs the production audit log repository
|
||||
// backed by the shared GORM connection.
|
||||
func NewAuditLogRepository(db *gorm.DB) interfaces.AuditLogRepository {
|
||||
return &auditLogRepository{db: db}
|
||||
}
|
||||
|
||||
// Create inserts a single audit row. Validation is light — Action is
|
||||
// the only required field at the schema level; CreatedAt is filled by
|
||||
// the database default if zero. Service-layer Log() fills both before
|
||||
// calling here so this is mostly a pass-through.
|
||||
func (r *auditLogRepository) Create(ctx context.Context, entry *types.AuditLog) error {
|
||||
return r.db.WithContext(ctx).Create(entry).Error
|
||||
}
|
||||
|
||||
// auditLogListLimitMax is the hard ceiling regardless of caller input.
|
||||
// Keeps a misconfigured client from triggering an unbounded scan; the
|
||||
// frontend's default page size is 50 so this is comfortable headroom.
|
||||
const auditLogListLimitMax = 100
|
||||
|
||||
// List returns audit log rows for a tenant, newest-first, applying the
|
||||
// optional cursor and filters from AuditLogQuery. Both Action and
|
||||
// Outcome are exact-match — the filter set is small enough that we
|
||||
// don't need substring matching yet.
|
||||
func (r *auditLogRepository) List(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
q *interfaces.AuditLogQuery,
|
||||
) ([]*types.AuditLog, error) {
|
||||
limit := 50
|
||||
if q != nil && q.Limit > 0 {
|
||||
limit = q.Limit
|
||||
}
|
||||
if limit > auditLogListLimitMax {
|
||||
limit = auditLogListLimitMax
|
||||
}
|
||||
|
||||
tx := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ?", tenantID)
|
||||
if q != nil {
|
||||
if q.AfterID > 0 {
|
||||
tx = tx.Where("id < ?", q.AfterID)
|
||||
}
|
||||
if q.Action != "" {
|
||||
tx = tx.Where("action = ?", q.Action)
|
||||
}
|
||||
if q.Outcome != "" {
|
||||
tx = tx.Where("outcome = ?", q.Outcome)
|
||||
}
|
||||
if q.ActorUserID != "" {
|
||||
tx = tx.Where("actor_user_id = ?", q.ActorUserID)
|
||||
}
|
||||
if q.ScopeType != "" {
|
||||
tx = tx.Where("scope_type = ?", q.ScopeType)
|
||||
}
|
||||
if q.ScopeID != "" {
|
||||
tx = tx.Where("scope_id = ?", q.ScopeID)
|
||||
}
|
||||
if q.UnscopedOnly {
|
||||
tx = tx.Where("scope_type = ''")
|
||||
}
|
||||
}
|
||||
|
||||
var entries []*types.AuditLog
|
||||
if err := tx.Order("id DESC").Limit(limit).Find(&entries).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return entries, nil
|
||||
}
|
||||
|
||||
// CountSinceForDedup powers AuditLogService.LogDenied's sliding-window
|
||||
// dedup. The (tenant_id, action) index makes the count cheap; the
|
||||
// remaining filters are applied as additional WHERE clauses on the
|
||||
// indexed result set.
|
||||
//
|
||||
// Returning early on the first match would be slightly cheaper, but
|
||||
// gorm's Count works on the same plan and the constant overhead is
|
||||
// negligible compared to the dedup benefit.
|
||||
func (r *auditLogRepository) CountSinceForDedup(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
actorUserID string,
|
||||
action types.AuditAction,
|
||||
requestPath string,
|
||||
since time.Time,
|
||||
) (int64, error) {
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).
|
||||
Model(&types.AuditLog{}).
|
||||
Where("tenant_id = ?", tenantID).
|
||||
Where("actor_user_id = ?", actorUserID).
|
||||
Where("action = ?", action).
|
||||
Where("request_path = ?", requestPath).
|
||||
Where("created_at >= ?", since).
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// DeleteOlderThan purges rows strictly older than cutoff in a single
|
||||
// DELETE. The retention sweep (driven by the audit log service) calls
|
||||
// it once a day with cutoff = now - retention_days.
|
||||
//
|
||||
// Tenant scope is intentionally not part of this signature: retention
|
||||
// is a global ops policy, not a per-tenant choice. If we ever need
|
||||
// per-tenant retention, we'd add a separate DeleteOlderThanForTenant
|
||||
// rather than overload this primitive.
|
||||
//
|
||||
// Returns the number of rows affected so the caller can log the sweep
|
||||
// outcome at INFO. Errors propagate verbatim — the caller decides
|
||||
// whether they're terminal or transient.
|
||||
func (r *auditLogRepository) DeleteOlderThan(ctx context.Context, cutoff time.Time) (int64, error) {
|
||||
res := r.db.WithContext(ctx).
|
||||
Where("created_at < ?", cutoff).
|
||||
Delete(&types.AuditLog{})
|
||||
if res.Error != nil {
|
||||
return 0, res.Error
|
||||
}
|
||||
return res.RowsAffected, nil
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestAuditLogRepositoryListFiltersKnowledgeBaseScope(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:audit-log-scope?mode=memory&cache=shared"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&types.AuditLog{}); err != nil {
|
||||
t.Fatalf("migrate audit log: %v", err)
|
||||
}
|
||||
rows := []*types.AuditLog{
|
||||
{TenantID: 7, Action: types.AuditActionMemberAdded},
|
||||
{TenantID: 7, Action: types.AuditActionKBUpdated, ScopeType: "knowledge_base", ScopeID: "kb-a"},
|
||||
{TenantID: 7, Action: types.AuditActionKnowledgeCreated, ScopeType: "knowledge_base", ScopeID: "kb-b"},
|
||||
{TenantID: 8, Action: types.AuditActionKBUpdated, ScopeType: "knowledge_base", ScopeID: "kb-a"},
|
||||
}
|
||||
for _, row := range rows {
|
||||
if err := db.Create(row).Error; err != nil {
|
||||
t.Fatalf("insert audit log: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
repo := NewAuditLogRepository(db)
|
||||
got, err := repo.List(context.Background(), 7, &interfaces.AuditLogQuery{
|
||||
ScopeType: "knowledge_base",
|
||||
ScopeID: "kb-a",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("list audit logs: %v", err)
|
||||
}
|
||||
if len(got) != 1 || got[0].TenantID != 7 || got[0].ScopeID != "kb-a" {
|
||||
t.Fatalf("scope filter returned unexpected rows: %+v", got)
|
||||
}
|
||||
unscoped, err := repo.List(context.Background(), 7, &interfaces.AuditLogQuery{UnscopedOnly: true})
|
||||
if err != nil {
|
||||
t.Fatalf("list unscoped audit logs: %v", err)
|
||||
}
|
||||
if len(unscoped) != 1 || unscoped[0].Action != types.AuditActionMemberAdded {
|
||||
t.Fatalf("unscoped filter returned unexpected rows: %+v", unscoped)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,117 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestDiffFAQChunkIDsByContentHash(t *testing.T) {
|
||||
toAdd, toDelete, matched := diffFAQChunkIDsByContentHash(
|
||||
[]chunkIDHash{
|
||||
{ID: "src-1", ContentHash: "hash-a"},
|
||||
{ID: "src-2", ContentHash: "hash-b"},
|
||||
{ID: "src-3", ContentHash: "hash-only-src"},
|
||||
},
|
||||
[]chunkIDHash{
|
||||
{ID: "dst-1", ContentHash: "hash-a"},
|
||||
{ID: "dst-2", ContentHash: "hash-only-dst"},
|
||||
},
|
||||
)
|
||||
|
||||
assert.ElementsMatch(t, []string{"src-2", "src-3"}, toAdd)
|
||||
assert.ElementsMatch(t, []string{"dst-2"}, toDelete)
|
||||
require.Len(t, matched, 1)
|
||||
assert.Equal(t, "src-1", matched[0].SrcChunkID)
|
||||
assert.Equal(t, "dst-1", matched[0].DstChunkID)
|
||||
}
|
||||
|
||||
func TestDiffFAQChunkIDsByContentHash_dstDuplicateHash(t *testing.T) {
|
||||
// Two destination chunks share hash-a (same FAQ content, e.g. stale copy after
|
||||
// answer_strategy-only change). Source has one chunk for hash-a → keep dst-1,
|
||||
// delete dst-dup, match for status sync.
|
||||
toAdd, toDelete, matched := diffFAQChunkIDsByContentHash(
|
||||
[]chunkIDHash{{ID: "src-1", ContentHash: "hash-a"}},
|
||||
[]chunkIDHash{
|
||||
{ID: "dst-1", ContentHash: "hash-a"},
|
||||
{ID: "dst-dup", ContentHash: "hash-a"},
|
||||
},
|
||||
)
|
||||
|
||||
assert.Empty(t, toAdd)
|
||||
assert.ElementsMatch(t, []string{"dst-dup"}, toDelete)
|
||||
require.Len(t, matched, 1)
|
||||
assert.Equal(t, "src-1", matched[0].SrcChunkID)
|
||||
assert.Equal(t, "dst-1", matched[0].DstChunkID)
|
||||
}
|
||||
|
||||
func TestDiffFAQChunkIDsByContentHash_dstDuplicateHash_notInSrc(t *testing.T) {
|
||||
// Duplicates whose hash is absent from source are all removed (orphan content).
|
||||
toAdd, toDelete, matched := diffFAQChunkIDsByContentHash(
|
||||
nil,
|
||||
[]chunkIDHash{
|
||||
{ID: "dst-orphan-1", ContentHash: "hash-orphan"},
|
||||
{ID: "dst-orphan-2", ContentHash: "hash-orphan"},
|
||||
},
|
||||
)
|
||||
|
||||
assert.Empty(t, toAdd)
|
||||
assert.ElementsMatch(t, []string{"dst-orphan-1", "dst-orphan-2"}, toDelete)
|
||||
assert.Empty(t, matched)
|
||||
}
|
||||
|
||||
func TestDiffFAQChunkIDsByContentHash_emptySide(t *testing.T) {
|
||||
toAdd, toDelete, matched := diffFAQChunkIDsByContentHash(
|
||||
[]chunkIDHash{{ID: "src-1", ContentHash: "hash-a"}},
|
||||
nil,
|
||||
)
|
||||
assert.Equal(t, []string{"src-1"}, toAdd)
|
||||
assert.Empty(t, toDelete)
|
||||
assert.Empty(t, matched)
|
||||
|
||||
toAdd, toDelete, matched = diffFAQChunkIDsByContentHash(
|
||||
nil,
|
||||
[]chunkIDHash{{ID: "dst-1", ContentHash: "hash-a"}},
|
||||
)
|
||||
assert.Empty(t, toAdd)
|
||||
assert.Equal(t, []string{"dst-1"}, toDelete)
|
||||
assert.Empty(t, matched)
|
||||
}
|
||||
|
||||
func TestFAQChunkDiff_SQLite(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
srcKBID := uuid.New().String()
|
||||
dstKBID := uuid.New().String()
|
||||
srcKnowledgeID := uuid.New().String()
|
||||
dstKnowledgeID := uuid.New().String()
|
||||
|
||||
require.NoError(t, db.Create(&types.Chunk{
|
||||
ID: uuid.New().String(), TenantID: 1, KnowledgeBaseID: srcKBID, KnowledgeID: srcKnowledgeID,
|
||||
ChunkType: types.ChunkTypeFAQ, ContentHash: "shared", Status: int(types.ChunkStatusIndexed),
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&types.Chunk{
|
||||
ID: uuid.New().String(), TenantID: 1, KnowledgeBaseID: srcKBID, KnowledgeID: srcKnowledgeID,
|
||||
ChunkType: types.ChunkTypeFAQ, ContentHash: "src-only", Status: int(types.ChunkStatusIndexed),
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&types.Chunk{
|
||||
ID: uuid.New().String(), TenantID: 1, KnowledgeBaseID: dstKBID, KnowledgeID: dstKnowledgeID,
|
||||
ChunkType: types.ChunkTypeFAQ, ContentHash: "shared", Status: int(types.ChunkStatusIndexed),
|
||||
}).Error)
|
||||
require.NoError(t, db.Create(&types.Chunk{
|
||||
ID: uuid.New().String(), TenantID: 1, KnowledgeBaseID: dstKBID, KnowledgeID: dstKnowledgeID,
|
||||
ChunkType: types.ChunkTypeFAQ, ContentHash: "dst-only", Status: int(types.ChunkStatusIndexed),
|
||||
}).Error)
|
||||
|
||||
diff, err := repo.FAQChunkDiff(ctx, 1, srcKBID, 1, dstKBID)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, diff.ChunksToAdd, 1)
|
||||
assert.Len(t, diff.ChunksToDelete, 1)
|
||||
assert.Len(t, diff.MatchedPairs, 1)
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestSaveChunkRevisionIsAtomicAndOptimistic(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:"+uuid.NewString()+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&types.Chunk{}, &types.ChunkRevision{}))
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
now := time.Now()
|
||||
chunk := &types.Chunk{
|
||||
ID: uuid.NewString(), TenantID: 1, KnowledgeBaseID: "kb", KnowledgeID: "knowledge",
|
||||
Content: "before", SourceContent: "before", ChunkType: types.ChunkTypeText,
|
||||
IsEnabled: true, IndexStatus: "ready", CreatedAt: now, UpdatedAt: now,
|
||||
}
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{chunk}))
|
||||
|
||||
snapshot := &types.ChunkRevision{
|
||||
ID: uuid.NewString(), TenantID: 1, KnowledgeBaseID: "kb", KnowledgeID: "knowledge",
|
||||
ChunkID: chunk.ID, Revision: 0, Content: "before", IsEnabled: true,
|
||||
EditSource: "user", EditedAt: now, CreatedAt: now,
|
||||
}
|
||||
chunk.Content = "after"
|
||||
chunk.ContentRevision = 1
|
||||
require.NoError(t, repo.SaveChunkRevision(ctx, chunk, snapshot, 0))
|
||||
|
||||
stored, err := repo.GetChunkByID(ctx, 1, chunk.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "after", stored.Content)
|
||||
require.Equal(t, 1, stored.ContentRevision)
|
||||
revisions, err := repo.ListChunkRevisions(ctx, 1, chunk.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, revisions, 1)
|
||||
require.Equal(t, "before", revisions[0].Content)
|
||||
|
||||
stale := *chunk
|
||||
stale.Content = "stale write"
|
||||
stale.ContentRevision = 1
|
||||
staleSnapshot := *snapshot
|
||||
staleSnapshot.ID = uuid.NewString()
|
||||
require.ErrorIs(t, repo.SaveChunkRevision(ctx, &stale, &staleSnapshot, 0), ErrChunkRevisionConflict)
|
||||
|
||||
stored, err = repo.GetChunkByID(ctx, 1, chunk.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "after", stored.Content)
|
||||
count := int64(0)
|
||||
require.NoError(t, db.Model(&types.ChunkRevision{}).Count(&count).Error)
|
||||
require.Equal(t, int64(1), count)
|
||||
require.False(t, errors.Is(gorm.ErrRecordNotFound, ErrChunkRevisionConflict))
|
||||
}
|
||||
@@ -0,0 +1,353 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// setupChunkTestDB creates an in-memory SQLite database with chunk and tag tables.
|
||||
func setupChunkTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&types.Chunk{}, &types.KnowledgeTag{}))
|
||||
return db
|
||||
}
|
||||
|
||||
func makeChunk(kbID, knowledgeID string, chunkType string) *types.Chunk {
|
||||
return &types.Chunk{
|
||||
ID: uuid.New().String(),
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: kbID,
|
||||
KnowledgeID: knowledgeID,
|
||||
Content: "test content",
|
||||
ChunkType: chunkType,
|
||||
IsEnabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChunks_SQLite_SeqIDAutoAssigned(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
kbID := uuid.New().String()
|
||||
knowledgeID := uuid.New().String()
|
||||
|
||||
// Create a batch of 5 chunks
|
||||
chunks := []*types.Chunk{
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
}
|
||||
|
||||
err := repo.CreateChunks(ctx, chunks)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify all chunks got unique sequential seq_ids
|
||||
var saved []types.Chunk
|
||||
require.NoError(t, db.Order("seq_id").Find(&saved).Error)
|
||||
assert.Len(t, saved, 5)
|
||||
|
||||
for i, c := range saved {
|
||||
assert.Equal(t, int64(i+1), c.SeqID, "chunk %d should have seq_id %d", i, i+1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChunks_CleansContextHeaderBeforePersistence(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
chunk := makeChunk(uuid.New().String(), uuid.New().String(), types.ChunkTypeText)
|
||||
chunk.ContextHeader = "# \x00投资评级说明\xff"
|
||||
|
||||
require.NoError(t, repo.CreateChunks(context.Background(), []*types.Chunk{chunk}))
|
||||
|
||||
require.Equal(t, "# 投资评级说明", chunk.ContextHeader)
|
||||
require.NotContains(t, chunk.ContextHeader, "\x00")
|
||||
require.True(t, utf8.ValidString(chunk.ContextHeader))
|
||||
|
||||
var saved types.Chunk
|
||||
require.NoError(t, db.First(&saved, "id = ?", chunk.ID).Error)
|
||||
assert.Equal(t, "# 投资评级说明", saved.ContextHeader)
|
||||
assert.NotContains(t, saved.ContextHeader, "\x00")
|
||||
assert.True(t, utf8.ValidString(saved.ContextHeader))
|
||||
}
|
||||
|
||||
func TestCreateChunks_SQLite_SeqIDContinuesFromExisting(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
kbID := uuid.New().String()
|
||||
knowledgeID := uuid.New().String()
|
||||
|
||||
// Create first batch
|
||||
batch1 := []*types.Chunk{
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
}
|
||||
require.NoError(t, repo.CreateChunks(ctx, batch1))
|
||||
|
||||
// Create second batch - seq_ids should continue from 3
|
||||
batch2 := []*types.Chunk{
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
}
|
||||
require.NoError(t, repo.CreateChunks(ctx, batch2))
|
||||
|
||||
var saved []types.Chunk
|
||||
require.NoError(t, db.Order("seq_id").Find(&saved).Error)
|
||||
assert.Len(t, saved, 5)
|
||||
|
||||
for i, c := range saved {
|
||||
assert.Equal(t, int64(i+1), c.SeqID, "chunk %d should have seq_id %d", i, i+1)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateChunks_SQLite_SeqIDUniqueAcrossKBs(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
kb1 := uuid.New().String()
|
||||
kb2 := uuid.New().String()
|
||||
k1 := uuid.New().String()
|
||||
k2 := uuid.New().String()
|
||||
|
||||
// Create chunks in two different knowledge bases
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{
|
||||
makeChunk(kb1, k1, "faq"),
|
||||
makeChunk(kb1, k1, "faq"),
|
||||
}))
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{
|
||||
makeChunk(kb2, k2, "faq"),
|
||||
makeChunk(kb2, k2, "faq"),
|
||||
}))
|
||||
|
||||
// All seq_ids should be globally unique (1,2,3,4)
|
||||
var saved []types.Chunk
|
||||
require.NoError(t, db.Order("seq_id").Find(&saved).Error)
|
||||
assert.Len(t, saved, 4)
|
||||
|
||||
seqIDs := map[int64]bool{}
|
||||
for _, c := range saved {
|
||||
assert.NotZero(t, c.SeqID)
|
||||
assert.False(t, seqIDs[c.SeqID], "seq_id %d should be unique", c.SeqID)
|
||||
seqIDs[c.SeqID] = true
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledgeTag_SQLite_SeqIDAutoAssigned(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
kbID := uuid.New().String()
|
||||
|
||||
// Create tags one by one (as the application does)
|
||||
tag1 := &types.KnowledgeTag{
|
||||
ID: uuid.New().String(),
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: kbID,
|
||||
Name: "tag1",
|
||||
}
|
||||
tag2 := &types.KnowledgeTag{
|
||||
ID: uuid.New().String(),
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: kbID,
|
||||
Name: "tag2",
|
||||
}
|
||||
|
||||
require.NoError(t, db.WithContext(ctx).Create(tag1).Error)
|
||||
require.NoError(t, db.WithContext(ctx).Create(tag2).Error)
|
||||
|
||||
// Both should have non-zero, unique seq_ids
|
||||
assert.NotZero(t, tag1.SeqID)
|
||||
assert.NotZero(t, tag2.SeqID)
|
||||
assert.NotEqual(t, tag1.SeqID, tag2.SeqID)
|
||||
}
|
||||
|
||||
func TestCreateChunks_SQLite_SeqIDAfterSoftDelete(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
kbID := uuid.New().String()
|
||||
knowledgeID := uuid.New().String()
|
||||
|
||||
// Create first batch
|
||||
batch1 := []*types.Chunk{
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
}
|
||||
require.NoError(t, repo.CreateChunks(ctx, batch1))
|
||||
|
||||
// Soft-delete all chunks (like frontend "clear" does)
|
||||
require.NoError(t, db.Where("knowledge_base_id = ?", kbID).Delete(&types.Chunk{}).Error)
|
||||
|
||||
// Verify soft-deleted
|
||||
var activeCount int64
|
||||
db.Model(&types.Chunk{}).Where("knowledge_base_id = ?", kbID).Count(&activeCount)
|
||||
assert.Equal(t, int64(0), activeCount, "all chunks should be soft-deleted")
|
||||
|
||||
// Create second batch — seq_ids must NOT conflict with soft-deleted ones
|
||||
batch2 := []*types.Chunk{
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
makeChunk(kbID, knowledgeID, "faq"),
|
||||
}
|
||||
err := repo.CreateChunks(ctx, batch2)
|
||||
require.NoError(t, err, "should not get UNIQUE constraint error after soft delete")
|
||||
|
||||
// Verify new seq_ids start after the soft-deleted max (3)
|
||||
var saved []types.Chunk
|
||||
require.NoError(t, db.Order("seq_id").Find(&saved).Error)
|
||||
assert.Len(t, saved, 2)
|
||||
assert.Equal(t, int64(4), saved[0].SeqID)
|
||||
assert.Equal(t, int64(5), saved[1].SeqID)
|
||||
}
|
||||
|
||||
func TestUpdateChunk_SQLite_NoNOWError(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
kbID := uuid.New().String()
|
||||
knowledgeID := uuid.New().String()
|
||||
|
||||
chunk := makeChunk(kbID, knowledgeID, "faq")
|
||||
require.NoError(t, db.WithContext(ctx).Create(chunk).Error)
|
||||
|
||||
// Test updating a chunk field — verifies no NOW() related errors
|
||||
err := db.WithContext(ctx).Model(chunk).Update("content", "updated content").Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
var saved types.Chunk
|
||||
require.NoError(t, db.First(&saved, "id = ?", chunk.ID).Error)
|
||||
assert.Equal(t, "updated content", saved.Content)
|
||||
}
|
||||
|
||||
func TestListPagedChunksByKnowledgeID_FiltersEnabledState(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
enabledChunk := makeChunk("kb-1", "faq-knowledge", types.ChunkTypeFAQ)
|
||||
disabledChunk := makeChunk("kb-1", "faq-knowledge", types.ChunkTypeFAQ)
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{enabledChunk, disabledChunk}))
|
||||
require.NoError(t, db.Model(&types.Chunk{}).
|
||||
Where("id = ?", disabledChunk.ID).
|
||||
Update("is_enabled", false).Error)
|
||||
|
||||
enabled := true
|
||||
chunks, total, err := repo.ListPagedChunksByKnowledgeID(
|
||||
ctx, 1, "faq-knowledge", &types.Pagination{Page: 1, PageSize: 20},
|
||||
[]types.ChunkType{types.ChunkTypeFAQ}, nil, "", "", "", types.KnowledgeTypeFAQ, &enabled,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), total)
|
||||
require.Len(t, chunks, 1)
|
||||
assert.Equal(t, enabledChunk.ID, chunks[0].ID)
|
||||
|
||||
allChunks, allTotal, err := repo.ListPagedChunksByKnowledgeID(
|
||||
ctx, 1, "faq-knowledge", &types.Pagination{Page: 1, PageSize: 20},
|
||||
[]types.ChunkType{types.ChunkTypeFAQ}, nil, "", "", "", types.KnowledgeTypeFAQ, nil,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), allTotal)
|
||||
assert.Len(t, allChunks, 2)
|
||||
}
|
||||
|
||||
func makeSuggestedFAQChunk(t *testing.T, kbID, knowledgeID, tagID, question string) *types.Chunk {
|
||||
t.Helper()
|
||||
chunk := makeChunk(kbID, knowledgeID, types.ChunkTypeFAQ)
|
||||
chunk.TagID = tagID
|
||||
chunk.Flags = types.ChunkFlagRecommended
|
||||
require.NoError(t, chunk.SetFAQMetadata(&types.FAQChunkMetadata{StandardQuestion: question}))
|
||||
return chunk
|
||||
}
|
||||
|
||||
func makeSuggestedDocumentChunk(t *testing.T, kbID, knowledgeID, question string) *types.Chunk {
|
||||
t.Helper()
|
||||
chunk := makeChunk(kbID, knowledgeID, types.ChunkTypeText)
|
||||
require.NoError(t, chunk.SetDocumentMetadata(&types.DocumentChunkMetadata{
|
||||
GeneratedQuestions: []types.GeneratedQuestion{{ID: uuid.NewString(), Question: question}},
|
||||
}))
|
||||
return chunk
|
||||
}
|
||||
|
||||
func TestListRecommendedFAQChunks_FiltersByTagWithoutWideningToParentKB(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
selectedTag := uuid.NewString()
|
||||
otherTag := uuid.NewString()
|
||||
selected := makeSuggestedFAQChunk(t, "kb-1", "faq-knowledge", selectedTag, "selected question")
|
||||
other := makeSuggestedFAQChunk(t, "kb-1", "faq-knowledge", otherTag, "other question")
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{selected, other}))
|
||||
|
||||
got, err := repo.ListRecommendedFAQChunks(ctx, 1, nil, nil, []string{selectedTag}, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 1)
|
||||
assert.Equal(t, selected.ID, got[0].ID)
|
||||
}
|
||||
|
||||
func TestListRecommendedFAQChunks_UnionsOnlyExplicitScopes(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
selectedTag := uuid.NewString()
|
||||
tagged := makeSuggestedFAQChunk(t, "kb-tag", "faq-tag", selectedTag, "tagged question")
|
||||
explicitKB := makeSuggestedFAQChunk(t, "kb-explicit", "faq-explicit", uuid.NewString(), "explicit KB question")
|
||||
unselected := makeSuggestedFAQChunk(t, "kb-other", "faq-other", uuid.NewString(), "unselected question")
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{tagged, explicitKB, unselected}))
|
||||
|
||||
got, err := repo.ListRecommendedFAQChunks(ctx, 1, []string{"kb-explicit"}, nil, []string{selectedTag}, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 2)
|
||||
assert.ElementsMatch(t, []string{tagged.ID, explicitKB.ID}, []string{got[0].ID, got[1].ID})
|
||||
}
|
||||
|
||||
func TestListRecentDocumentChunksWithQuestions_KnowledgeScopeDoesNotIncludeSiblingDocuments(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
selected := makeSuggestedDocumentChunk(t, "kb-1", "doc-selected", "selected document question")
|
||||
sibling := makeSuggestedDocumentChunk(t, "kb-1", "doc-sibling", "sibling document question")
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{selected, sibling}))
|
||||
|
||||
got, err := repo.ListRecentDocumentChunksWithQuestions(ctx, 1, nil, []string{"doc-selected"}, 10)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 1)
|
||||
assert.Equal(t, selected.ID, got[0].ID)
|
||||
}
|
||||
|
||||
func TestListRecentDocumentChunksWithQuestions_UnionsExplicitKBAndKnowledge(t *testing.T) {
|
||||
db := setupChunkTestDB(t)
|
||||
repo := NewChunkRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
fromExplicitKB := makeSuggestedDocumentChunk(t, "kb-explicit", "doc-1", "explicit KB question")
|
||||
fromExplicitDocument := makeSuggestedDocumentChunk(t, "kb-other", "doc-selected", "selected document question")
|
||||
unselected := makeSuggestedDocumentChunk(t, "kb-other", "doc-other", "unselected question")
|
||||
require.NoError(t, repo.CreateChunks(ctx, []*types.Chunk{fromExplicitKB, fromExplicitDocument, unselected}))
|
||||
|
||||
got, err := repo.ListRecentDocumentChunksWithQuestions(
|
||||
ctx, 1, []string{"kb-explicit"}, []string{"doc-selected"}, 10,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, got, 2)
|
||||
assert.ElementsMatch(t, []string{fromExplicitKB.ID, fromExplicitDocument.ID}, []string{got[0].ID, got[1].ID})
|
||||
}
|
||||
@@ -0,0 +1,139 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ErrCustomAgentNotFound is returned when a custom agent is not found
|
||||
var ErrCustomAgentNotFound = errors.New("custom agent not found")
|
||||
|
||||
// customAgentRepository implements the CustomAgentRepository interface
|
||||
type customAgentRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewCustomAgentRepository creates a new custom agent repository
|
||||
func NewCustomAgentRepository(db *gorm.DB) interfaces.CustomAgentRepository {
|
||||
return &customAgentRepository{db: db}
|
||||
}
|
||||
|
||||
// CreateAgent creates a new custom agent
|
||||
func (r *customAgentRepository) CreateAgent(ctx context.Context, agent *types.CustomAgent) error {
|
||||
return r.db.WithContext(ctx).Create(agent).Error
|
||||
}
|
||||
|
||||
// GetAgentByID gets an agent by id and tenant
|
||||
func (r *customAgentRepository) GetAgentByID(ctx context.Context, id string, tenantID uint64) (*types.CustomAgent, error) {
|
||||
var agent types.CustomAgent
|
||||
if err := r.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&agent).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrCustomAgentNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &agent, nil
|
||||
}
|
||||
|
||||
// ListAgentsByTenantID lists all agents for a specific tenant
|
||||
func (r *customAgentRepository) ListAgentsByTenantID(ctx context.Context, tenantID uint64) ([]*types.CustomAgent, error) {
|
||||
var agents []*types.CustomAgent
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ?", tenantID).
|
||||
Order("created_at DESC").
|
||||
Find(&agents).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return agents, nil
|
||||
}
|
||||
|
||||
// UpdateAgent updates an agent
|
||||
func (r *customAgentRepository) UpdateAgent(ctx context.Context, agent *types.CustomAgent) error {
|
||||
return r.db.WithContext(ctx).Save(agent).Error
|
||||
}
|
||||
|
||||
// DeleteAgent deletes an agent (soft delete)
|
||||
func (r *customAgentRepository) DeleteAgent(ctx context.Context, id string, tenantID uint64) error {
|
||||
return r.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&types.CustomAgent{}).Error
|
||||
}
|
||||
|
||||
// CountByModelID counts active agents whose config references modelID.
|
||||
func (r *customAgentRepository) CountByModelID(
|
||||
ctx context.Context, tenantID uint64, modelID string,
|
||||
) (int64, error) {
|
||||
var count int64
|
||||
query := r.db.WithContext(ctx).
|
||||
Model(&types.CustomAgent{}).
|
||||
Where("tenant_id = ?", tenantID)
|
||||
query = scopeCustomAgentsByModelID(query, modelID)
|
||||
err := query.Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// ListModelUsages returns active agents that reference modelID. Only id, name,
|
||||
// and config are loaded; the config is decoded locally and never returned.
|
||||
func (r *customAgentRepository) ListModelUsages(
|
||||
ctx context.Context, tenantID uint64, modelID string,
|
||||
) ([]types.ModelUsageResource, error) {
|
||||
rows := make([]*types.CustomAgent, 0)
|
||||
query := r.db.WithContext(ctx).
|
||||
Model(&types.CustomAgent{}).
|
||||
Select("id", "name", "config").
|
||||
Where("tenant_id = ?", tenantID)
|
||||
query = scopeCustomAgentsByModelID(query, modelID)
|
||||
if err := query.Order("name ASC, id ASC").Limit(types.ModelUsageListLimit).Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
usages := make([]types.ModelUsageResource, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
bindings := customAgentModelUsageBindings(row, modelID)
|
||||
if len(bindings) == 0 {
|
||||
continue
|
||||
}
|
||||
usages = append(usages, types.ModelUsageResource{
|
||||
ID: row.ID,
|
||||
Name: row.Name,
|
||||
Bindings: bindings,
|
||||
})
|
||||
}
|
||||
return usages, nil
|
||||
}
|
||||
|
||||
// CountBySandboxConfigID counts agents pointing at a sandbox config.
|
||||
//
|
||||
// Used only to warn the admin which agents reference a config; never use it to
|
||||
// refuse operations. Agent references are permanent state, so blocking on them
|
||||
// would make credential rotation impossible.
|
||||
func (r *customAgentRepository) CountBySandboxConfigID(
|
||||
ctx context.Context, tenantID uint64, configID string,
|
||||
) (int64, error) {
|
||||
var count int64
|
||||
query := r.db.WithContext(ctx).
|
||||
Model(&types.CustomAgent{}).
|
||||
Where("tenant_id = ?", tenantID)
|
||||
query = scopeCustomAgentsBySandboxConfigID(query, configID)
|
||||
err := query.Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// ListNamesBySandboxConfigID returns agent names pointing at a sandbox config.
|
||||
//
|
||||
// Used only to warn the admin which agents reference a config; never use it to
|
||||
// refuse operations. Agent references are permanent state, so blocking on them
|
||||
// would make credential rotation impossible.
|
||||
func (r *customAgentRepository) ListNamesBySandboxConfigID(
|
||||
ctx context.Context, tenantID uint64, configID string,
|
||||
) ([]string, error) {
|
||||
var names []string
|
||||
query := r.db.WithContext(ctx).
|
||||
Model(&types.CustomAgent{}).
|
||||
Where("tenant_id = ?", tenantID)
|
||||
query = scopeCustomAgentsBySandboxConfigID(query, configID)
|
||||
err := query.Order("name ASC").Pluck("name", &names).Error
|
||||
return names, err
|
||||
}
|
||||
@@ -0,0 +1,324 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// DataSourceRepository provides data access for data sources
|
||||
type DataSourceRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewDataSourceRepository creates a new data source repository
|
||||
func NewDataSourceRepository(db *gorm.DB) interfaces.DataSourceRepository {
|
||||
return &DataSourceRepository{db: db}
|
||||
}
|
||||
|
||||
// Create inserts a new data source record
|
||||
func (r *DataSourceRepository) Create(ctx context.Context, ds *types.DataSource) error {
|
||||
if ds == nil {
|
||||
return errors.New("data source is nil")
|
||||
}
|
||||
// GORM treats false as the zero value of bool. For a field tagged
|
||||
// default:true it replaces both the INSERT value and the in-memory field
|
||||
// with true, so a caller-selected false would be lost. Capture it, force
|
||||
// the column write, then restore the struct so Create's return value (and
|
||||
// the HTTP 201 body) match the database.
|
||||
syncDeletions := ds.SyncDeletions
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(ds).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Model(&types.DataSource{}).
|
||||
Where("id = ?", ds.ID).
|
||||
UpdateColumn("sync_deletions", syncDeletions).Error
|
||||
})
|
||||
ds.SyncDeletions = syncDeletions
|
||||
return err
|
||||
}
|
||||
|
||||
// FindByID retrieves a data source by ID
|
||||
func (r *DataSourceRepository) FindByID(ctx context.Context, id string) (*types.DataSource, error) {
|
||||
if id == "" {
|
||||
return nil, errors.New("id is empty")
|
||||
}
|
||||
var ds types.DataSource
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("id = ?", id).
|
||||
Where("deleted_at IS NULL").
|
||||
First(&ds).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, errors.New("data source not found")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &ds, nil
|
||||
}
|
||||
|
||||
// FindByKnowledgeBase lists all data sources for a knowledge base
|
||||
func (r *DataSourceRepository) FindByKnowledgeBase(ctx context.Context, kbID string) ([]*types.DataSource, error) {
|
||||
if kbID == "" {
|
||||
return nil, errors.New("knowledge base id is empty")
|
||||
}
|
||||
var dataSources []*types.DataSource
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("knowledge_base_id = ?", kbID).
|
||||
Where("deleted_at IS NULL").
|
||||
Order("created_at DESC").
|
||||
Find(&dataSources).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return dataSources, nil
|
||||
}
|
||||
|
||||
// Update updates an existing data source
|
||||
func (r *DataSourceRepository) Update(ctx context.Context, ds *types.DataSource) error {
|
||||
if ds == nil {
|
||||
return errors.New("data source is nil")
|
||||
}
|
||||
if ds.ID == "" {
|
||||
return errors.New("data source id is empty")
|
||||
}
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Model(ds).Updates(ds).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// GORM Updates(struct) deliberately skips zero values, which would make
|
||||
// a user-selected sync_deletions=false impossible to persist.
|
||||
return tx.Model(&types.DataSource{}).
|
||||
Where("id = ?", ds.ID).
|
||||
UpdateColumn("sync_deletions", ds.SyncDeletions).Error
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateSyncState updates only fields managed by sync execution. GORM's
|
||||
// Updates(struct) skips zero values, so use a map here to persist cleared error
|
||||
// messages without broadening the generic Update method.
|
||||
func (r *DataSourceRepository) UpdateSyncState(ctx context.Context, ds *types.DataSource) error {
|
||||
if ds == nil {
|
||||
return errors.New("data source is nil")
|
||||
}
|
||||
if ds.ID == "" {
|
||||
return errors.New("data source id is empty")
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&types.DataSource{}).
|
||||
Where("id = ?", ds.ID).
|
||||
Updates(map[string]interface{}{
|
||||
"status": ds.Status,
|
||||
"last_sync_at": ds.LastSyncAt,
|
||||
"last_sync_cursor": ds.LastSyncCursor,
|
||||
"last_sync_result": ds.LastSyncResult,
|
||||
"error_message": ds.ErrorMessage,
|
||||
"updated_at": time.Now().UTC(),
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Delete performs a soft delete
|
||||
func (r *DataSourceRepository) Delete(ctx context.Context, id string) error {
|
||||
if id == "" {
|
||||
return errors.New("id is empty")
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("id = ?", id).
|
||||
Delete(&types.DataSource{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// FindActive retrieves all active data sources (used for scheduling)
|
||||
func (r *DataSourceRepository) FindActive(ctx context.Context) ([]*types.DataSource, error) {
|
||||
var dataSources []*types.DataSource
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("status = ?", types.DataSourceStatusActive).
|
||||
Where("deleted_at IS NULL").
|
||||
Where("sync_schedule != ''").
|
||||
Order("created_at DESC").
|
||||
Find(&dataSources).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return dataSources, nil
|
||||
}
|
||||
|
||||
// SyncLogRepository provides data access for sync logs
|
||||
type SyncLogRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewSyncLogRepository creates a new sync log repository
|
||||
func NewSyncLogRepository(db *gorm.DB) interfaces.SyncLogRepository {
|
||||
return &SyncLogRepository{db: db}
|
||||
}
|
||||
|
||||
// Create inserts a new sync log entry
|
||||
func (r *SyncLogRepository) Create(ctx context.Context, log *types.SyncLog) error {
|
||||
if log == nil {
|
||||
return errors.New("sync log is nil")
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Create(log).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// FindByID retrieves a sync log by ID
|
||||
func (r *SyncLogRepository) FindByID(ctx context.Context, id string) (*types.SyncLog, error) {
|
||||
if id == "" {
|
||||
return nil, errors.New("id is empty")
|
||||
}
|
||||
var log types.SyncLog
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("id = ?", id).
|
||||
First(&log).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, errors.New("sync log not found")
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &log, nil
|
||||
}
|
||||
|
||||
// FindByDataSource lists sync logs for a data source with pagination
|
||||
func (r *SyncLogRepository) FindByDataSource(ctx context.Context, dsID string, limit int, offset int) ([]*types.SyncLog, error) {
|
||||
if dsID == "" {
|
||||
return nil, errors.New("data source id is empty")
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
var logs []*types.SyncLog
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("data_source_id = ?", dsID).
|
||||
Order("started_at DESC").
|
||||
Limit(limit).
|
||||
Offset(offset).
|
||||
Find(&logs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return logs, nil
|
||||
}
|
||||
|
||||
// FindLatest retrieves the most recent sync log for a data source
|
||||
func (r *SyncLogRepository) FindLatest(ctx context.Context, dsID string) (*types.SyncLog, error) {
|
||||
if dsID == "" {
|
||||
return nil, errors.New("data source id is empty")
|
||||
}
|
||||
var log types.SyncLog
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("data_source_id = ?", dsID).
|
||||
Order("started_at DESC").
|
||||
Limit(1).
|
||||
First(&log).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &log, nil
|
||||
}
|
||||
|
||||
// HasRunningSync checks if a data source has any sync currently in "running" status.
|
||||
func (r *SyncLogRepository) HasRunningSync(ctx context.Context, dsID string) (bool, error) {
|
||||
if dsID == "" {
|
||||
return false, errors.New("data source id is empty")
|
||||
}
|
||||
var count int64
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&types.SyncLog{}).
|
||||
Where("data_source_id = ?", dsID).
|
||||
Where("status = ?", types.SyncLogStatusRunning).
|
||||
Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
return count > 0, nil
|
||||
}
|
||||
|
||||
// Update updates an existing sync log entry
|
||||
func (r *SyncLogRepository) Update(ctx context.Context, log *types.SyncLog) error {
|
||||
if log == nil {
|
||||
return errors.New("sync log is nil")
|
||||
}
|
||||
if log.ID == "" {
|
||||
return errors.New("sync log id is empty")
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(log).
|
||||
Updates(log).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateResult updates only fields produced by sync execution. Use an explicit
|
||||
// map so empty error messages are written when a later sync succeeds.
|
||||
func (r *SyncLogRepository) UpdateResult(ctx context.Context, log *types.SyncLog) error {
|
||||
if log == nil {
|
||||
return errors.New("sync log is nil")
|
||||
}
|
||||
if log.ID == "" {
|
||||
return errors.New("sync log id is empty")
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&types.SyncLog{}).
|
||||
Where("id = ?", log.ID).
|
||||
Updates(map[string]interface{}{
|
||||
"status": log.Status,
|
||||
"finished_at": log.FinishedAt,
|
||||
"items_total": log.ItemsTotal,
|
||||
"items_created": log.ItemsCreated,
|
||||
"items_updated": log.ItemsUpdated,
|
||||
"items_deleted": log.ItemsDeleted,
|
||||
"items_skipped": log.ItemsSkipped,
|
||||
"items_failed": log.ItemsFailed,
|
||||
"error_message": log.ErrorMessage,
|
||||
"result": log.Result,
|
||||
"updated_at": time.Now().UTC(),
|
||||
}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CancelPendingByDataSource marks all non-terminal sync logs for a data source as canceled.
|
||||
func (r *SyncLogRepository) CancelPendingByDataSource(ctx context.Context, dsID string) error {
|
||||
if dsID == "" {
|
||||
return errors.New("data source id is empty")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.SyncLog{}).
|
||||
Where("data_source_id = ?", dsID).
|
||||
Where("status IN ?", []string{types.SyncLogStatusRunning, "pending"}).
|
||||
Updates(map[string]interface{}{
|
||||
"status": types.SyncLogStatusCanceled,
|
||||
"finished_at": &now,
|
||||
"error_message": "data source deleted",
|
||||
}).Error
|
||||
}
|
||||
|
||||
// CleanupOldLogs deletes sync logs older than the retention period
|
||||
func (r *SyncLogRepository) CleanupOldLogs(ctx context.Context, retentionDays int) error {
|
||||
if retentionDays <= 0 {
|
||||
retentionDays = 30
|
||||
}
|
||||
// Delete logs older than the retention period
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("started_at < NOW() - INTERVAL ? DAY", retentionDays).
|
||||
Delete(&types.SyncLog{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupDataSourceRepoTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&types.DataSource{}, &types.SyncLog{}))
|
||||
return db
|
||||
}
|
||||
|
||||
func TestDataSourceRepositoryUpdateSyncStateClearsErrorMessage(t *testing.T) {
|
||||
db := setupDataSourceRepoTestDB(t)
|
||||
repo := NewDataSourceRepository(db)
|
||||
now := time.Now().UTC()
|
||||
result := types.JSON(`{"total":0}`)
|
||||
|
||||
ds := &types.DataSource{
|
||||
ID: "ds-1",
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Name: "Feishu",
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
Status: types.DataSourceStatusError,
|
||||
ErrorMessage: "previous failure",
|
||||
}
|
||||
require.NoError(t, repo.Create(context.Background(), ds))
|
||||
|
||||
ds.Status = types.DataSourceStatusActive
|
||||
ds.ErrorMessage = ""
|
||||
ds.LastSyncAt = &now
|
||||
ds.LastSyncResult = result
|
||||
require.NoError(t, repo.UpdateSyncState(context.Background(), ds))
|
||||
|
||||
var stored types.DataSource
|
||||
require.NoError(t, db.First(&stored, "id = ?", ds.ID).Error)
|
||||
assert.Equal(t, types.DataSourceStatusActive, stored.Status)
|
||||
assert.Empty(t, stored.ErrorMessage)
|
||||
assert.Equal(t, result.ToString(), stored.LastSyncResult.ToString())
|
||||
require.NotNil(t, stored.LastSyncAt)
|
||||
}
|
||||
|
||||
func TestDataSourceRepositoryUpdatePersistsDisabledSyncDeletions(t *testing.T) {
|
||||
db := setupDataSourceRepoTestDB(t)
|
||||
repo := NewDataSourceRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
ds := &types.DataSource{
|
||||
ID: "ds-sync-deletions",
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Name: "Feishu",
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
SyncDeletions: true,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, ds))
|
||||
|
||||
ds.SyncDeletions = false
|
||||
require.NoError(t, repo.Update(ctx, ds))
|
||||
assert.False(t, ds.SyncDeletions)
|
||||
|
||||
var stored types.DataSource
|
||||
require.NoError(t, db.First(&stored, "id = ?", ds.ID).Error)
|
||||
assert.False(t, stored.SyncDeletions)
|
||||
|
||||
ds.SyncDeletions = true
|
||||
require.NoError(t, repo.Update(ctx, ds))
|
||||
assert.True(t, ds.SyncDeletions)
|
||||
require.NoError(t, db.First(&stored, "id = ?", ds.ID).Error)
|
||||
assert.True(t, stored.SyncDeletions)
|
||||
}
|
||||
|
||||
func TestDataSourceRepositoryCreatePersistsDisabledSyncDeletions(t *testing.T) {
|
||||
db := setupDataSourceRepoTestDB(t)
|
||||
repo := NewDataSourceRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
ds := &types.DataSource{
|
||||
ID: "ds-create-sync-deletions",
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Name: "Feishu",
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
SyncDeletions: false,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, ds))
|
||||
assert.False(t, ds.SyncDeletions, "Create must not leave the in-memory field hydrated to the GORM default")
|
||||
|
||||
var stored types.DataSource
|
||||
require.NoError(t, db.First(&stored, "id = ?", ds.ID).Error)
|
||||
assert.False(t, stored.SyncDeletions)
|
||||
|
||||
// A later Updates() on the same pointer must not persist the hydrated default.
|
||||
ds.Name = "Renamed"
|
||||
require.NoError(t, repo.Update(ctx, ds))
|
||||
require.NoError(t, db.First(&stored, "id = ?", ds.ID).Error)
|
||||
assert.False(t, stored.SyncDeletions)
|
||||
assert.Equal(t, "Renamed", stored.Name)
|
||||
}
|
||||
|
||||
func TestDataSourceRepositoryCreatePersistsEnabledSyncDeletions(t *testing.T) {
|
||||
db := setupDataSourceRepoTestDB(t)
|
||||
repo := NewDataSourceRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
ds := &types.DataSource{
|
||||
ID: "ds-create-sync-deletions-enabled",
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Name: "Feishu",
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
SyncDeletions: true,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, ds))
|
||||
assert.True(t, ds.SyncDeletions)
|
||||
|
||||
var stored types.DataSource
|
||||
require.NoError(t, db.First(&stored, "id = ?", ds.ID).Error)
|
||||
assert.True(t, stored.SyncDeletions)
|
||||
}
|
||||
|
||||
func TestDataSourceRepositoryDeleteSoftDeletesOnSQLite(t *testing.T) {
|
||||
db := setupDataSourceRepoTestDB(t)
|
||||
repo := NewDataSourceRepository(db)
|
||||
ctx := context.Background()
|
||||
|
||||
target := &types.DataSource{
|
||||
ID: "ds-delete-target",
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Name: "Delete target",
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
}
|
||||
other := &types.DataSource{
|
||||
ID: "ds-delete-other",
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Name: "Other data source",
|
||||
Type: types.ConnectorTypeFeishu,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, target))
|
||||
require.NoError(t, repo.Create(ctx, other))
|
||||
|
||||
require.NoError(t, repo.Delete(ctx, target.ID))
|
||||
|
||||
var deleted types.DataSource
|
||||
require.NoError(t, db.Unscoped().First(&deleted, "id = ?", target.ID).Error)
|
||||
assert.True(t, deleted.DeletedAt.Valid)
|
||||
|
||||
found, err := repo.FindByID(ctx, target.ID)
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, found)
|
||||
|
||||
untouched, err := repo.FindByID(ctx, other.ID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, other.ID, untouched.ID)
|
||||
}
|
||||
|
||||
func TestSyncLogRepositoryUpdateResultClearsErrorMessage(t *testing.T) {
|
||||
db := setupDataSourceRepoTestDB(t)
|
||||
repo := NewSyncLogRepository(db)
|
||||
finishedAt := time.Now().UTC()
|
||||
result := types.JSON(`{"total":0}`)
|
||||
|
||||
log := &types.SyncLog{
|
||||
ID: "log-1",
|
||||
DataSourceID: "ds-1",
|
||||
TenantID: 1,
|
||||
Status: types.SyncLogStatusFailed,
|
||||
ErrorMessage: "previous failure",
|
||||
ItemsTotal: 1,
|
||||
ItemsFailed: 1,
|
||||
}
|
||||
require.NoError(t, repo.Create(context.Background(), log))
|
||||
|
||||
log.Status = types.SyncLogStatusSuccess
|
||||
log.ErrorMessage = ""
|
||||
log.FinishedAt = &finishedAt
|
||||
log.ItemsTotal = 0
|
||||
log.ItemsFailed = 0
|
||||
log.Result = result
|
||||
require.NoError(t, repo.UpdateResult(context.Background(), log))
|
||||
|
||||
var stored types.SyncLog
|
||||
require.NoError(t, db.First(&stored, "id = ?", log.ID).Error)
|
||||
assert.Equal(t, types.SyncLogStatusSuccess, stored.Status)
|
||||
assert.Empty(t, stored.ErrorMessage)
|
||||
assert.Zero(t, stored.ItemsTotal)
|
||||
assert.Zero(t, stored.ItemsFailed)
|
||||
assert.Equal(t, result.ToString(), stored.Result.ToString())
|
||||
require.NotNil(t, stored.FinishedAt)
|
||||
}
|
||||
@@ -0,0 +1,567 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type DigitalHumanRepository struct{ db *gorm.DB }
|
||||
|
||||
func NewDigitalHumanRepository(db *gorm.DB) *DigitalHumanRepository {
|
||||
return &DigitalHumanRepository{db: db}
|
||||
}
|
||||
|
||||
func dhScope(db *gorm.DB, scope types.DigitalHumanScope) *gorm.DB {
|
||||
return db.Where("tenant_id = ? AND owner_id = ?", scope.TenantID, scope.OwnerID)
|
||||
}
|
||||
|
||||
func dhReadScope(db *gorm.DB, scope types.DigitalHumanScope) *gorm.DB {
|
||||
if scope.ReadAll {
|
||||
return db.Where("tenant_id = ?", scope.TenantID)
|
||||
}
|
||||
q := db.Where("tenant_id = ?", scope.TenantID).Where("owner_id = ? OR id IN ?", scope.OwnerID, accesspolicy.SharedIDs(db.Statement.Context, "digital-humans"))
|
||||
if ids := accesspolicy.HiddenIDs(db.Statement.Context, "digital-humans"); len(ids) > 0 {
|
||||
q = q.Where("id NOT IN ?", ids)
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) List(ctx context.Context, scope types.DigitalHumanScope, offset, limit int) ([]types.DigitalHuman, error) {
|
||||
result := []types.DigitalHuman{}
|
||||
err := dhReadScope(r.db.WithContext(ctx), scope).Order("created_at DESC, id ASC").Offset(offset).Limit(limit).Find(&result).Error
|
||||
if err == nil {
|
||||
humans := make([]*types.DigitalHuman, len(result))
|
||||
for i := range result {
|
||||
humans[i] = &result[i]
|
||||
}
|
||||
err = r.populateCreatorNames(ctx, humans)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) populateCreatorNames(ctx context.Context, humans []*types.DigitalHuman) error {
|
||||
if len(humans) == 0 {
|
||||
return nil
|
||||
}
|
||||
ids := make([]string, 0, len(humans))
|
||||
for _, human := range humans {
|
||||
ids = append(ids, human.OwnerID)
|
||||
}
|
||||
var users []struct{ ID, Username string }
|
||||
if err := r.db.WithContext(ctx).Model(&types.User{}).Select("id, username").Where("id IN ?", ids).Find(&users).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
names := make(map[string]string, len(users))
|
||||
for _, user := range users {
|
||||
names[user.ID] = user.Username
|
||||
}
|
||||
for _, human := range humans {
|
||||
human.CreatorName = names[human.OwnerID]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Get(ctx context.Context, scope types.DigitalHumanScope, id string) (*types.DigitalHumanDetail, error) {
|
||||
result := &types.DigitalHumanDetail{DigitalHuman: &types.DigitalHuman{}, Versions: []types.DigitalHumanVersion{}, Bindings: []types.DigitalHumanBinding{}}
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := dhReadScope(tx, scope).First(result.DigitalHuman, "id = ?", id).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
ownerScope := types.DigitalHumanScope{TenantID: result.TenantID, OwnerID: result.OwnerID}
|
||||
// Raw source documents are deliberately excluded from normal detail responses.
|
||||
if err := dhScope(tx, ownerScope).Omit("sources").Where("digital_human_id = ?", id).Order("number DESC").Limit(100).Find(&result.Versions).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
bindings := tx.Where("tenant_id = ? AND digital_human_id = ?", scope.TenantID, id)
|
||||
if !scope.ReadAll && scope.OwnerID != result.OwnerID {
|
||||
bindings = bindings.Where("owner_id = ?", scope.OwnerID)
|
||||
}
|
||||
if err := bindings.Find(&result.Bindings).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var jobs []types.DigitalHumanJob
|
||||
if err := dhScope(tx, ownerScope).Where("digital_human_id = ?", id).Order("created_at DESC, id ASC").Limit(1).Find(&jobs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(jobs) > 0 {
|
||||
result.LastJob = &jobs[0]
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err == nil {
|
||||
err = r.populateCreatorNames(ctx, []*types.DigitalHuman{result.DigitalHuman})
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func dhNotFound(err error) error {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return types.ErrDigitalHumanNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// All mutations lock the owner-scoped human first; related agents and versions
|
||||
// are committed in the same transaction, so review cannot partly publish.
|
||||
func (r *DigitalHumanRepository) mutate(ctx context.Context, scope types.DigitalHumanScope, id string, fn func(*gorm.DB, *types.DigitalHuman) error) error {
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var human types.DigitalHuman
|
||||
if err := dhScope(tx.Clauses(clause.Locking{Strength: "UPDATE"}), scope).First(&human, "id = ?", id).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
return fn(tx, &human)
|
||||
})
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Create(ctx context.Context, human *types.DigitalHuman, persona map[string]string) error {
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
base := "按用户当前任务回答,保证上下文衔接和事实准确。数字人表达档案须经人工审核后生效。"
|
||||
agent := &types.CustomAgent{ID: human.PrimaryAgentID, TenantID: human.TenantID, CreatedBy: human.OwnerID,
|
||||
Name: human.Name, Description: "数字人「" + human.Name + "」的默认对话 Agent", Avatar: human.Avatar,
|
||||
Config: types.CustomAgentConfig{AgentMode: types.AgentModeSmartReasoning, AgentType: types.AgentTypeCustom,
|
||||
ModelID: human.ModelID, SystemPrompt: base, KBSelectionMode: "none", SkillsSelectionMode: "none", MCPSelectionMode: "none"}}
|
||||
agent.EnsureDefaults()
|
||||
if err := tx.Create(agent).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Create(human).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
binding := &types.DigitalHumanBinding{ID: uuid.NewString(), TenantID: human.TenantID, OwnerID: human.OwnerID,
|
||||
DigitalHumanID: human.ID, AgentID: agent.ID, BasePrompt: base}
|
||||
if err := tx.Create(binding).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(&types.DigitalHumanVersion{ID: uuid.NewString(), TenantID: human.TenantID, OwnerID: human.OwnerID,
|
||||
DigitalHumanID: human.ID, Number: 1, Status: "draft", Persona: persona, Sources: []types.DigitalHumanSource{}}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Update(ctx context.Context, scope types.DigitalHumanScope, id string, revision int64, name, description, avatar, modelID string, persona map[string]string) error {
|
||||
return r.mutate(ctx, scope, id, func(tx *gorm.DB, human *types.DigitalHuman) error {
|
||||
if human.Revision != revision || human.ProcessingJobID != "" {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
if persona != nil {
|
||||
var previous types.DigitalHumanVersion
|
||||
if err := dhScope(tx, scope).First(&previous, "digital_human_id = ? AND number = ?", id, human.LatestVersion).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
human.LatestVersion++
|
||||
v := types.DigitalHumanVersion{ID: uuid.NewString(), TenantID: human.TenantID, OwnerID: human.OwnerID, DigitalHumanID: id,
|
||||
Number: human.LatestVersion, Status: "draft", Persona: persona, Sources: previous.Sources}
|
||||
if err := tx.Create(&v).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
human.Status = "testing"
|
||||
}
|
||||
human.Name, human.Description, human.Avatar, human.ModelID = name, description, avatar, modelID
|
||||
human.Revision++
|
||||
if err := tx.Save(human).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var agent types.CustomAgent
|
||||
if err := tx.First(&agent, "id = ? AND tenant_id = ? AND created_by = ?", human.PrimaryAgentID, scope.TenantID, scope.OwnerID).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
agent.Name, agent.Avatar, agent.Config.ModelID = name, avatar, modelID
|
||||
agent.Description = "数字人「" + name + "」的默认对话 Agent"
|
||||
return tx.Save(&agent).Error
|
||||
})
|
||||
}
|
||||
|
||||
func digitalHumanPromptMarkers(id string) (string, string) {
|
||||
return "<!-- weknora:digital-human:" + id + ":start -->", "<!-- weknora:digital-human:" + id + ":end -->"
|
||||
}
|
||||
|
||||
func removeDigitalHumanPrompt(prompt, id string) string {
|
||||
start, end := digitalHumanPromptMarkers(id)
|
||||
for {
|
||||
first := strings.Index(prompt, start)
|
||||
if first < 0 {
|
||||
return prompt
|
||||
}
|
||||
last := strings.Index(prompt[first+len(start):], end)
|
||||
if last < 0 {
|
||||
return prompt
|
||||
}
|
||||
last += first + len(start) + len(end)
|
||||
prompt = strings.TrimRight(prompt[:first], "\n") + prompt[last:]
|
||||
}
|
||||
}
|
||||
|
||||
func digitalHumanPrompt(human *types.DigitalHuman, version *types.DigitalHumanVersion, base string) string {
|
||||
keys := make([]string, 0, len(version.Persona))
|
||||
for key := range version.Persona {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
var b strings.Builder
|
||||
b.WriteString(removeDigitalHumanPrompt(base, human.ID))
|
||||
start, end := digitalHumanPromptMarkers(human.ID)
|
||||
label := "已审核数字人"
|
||||
if version.Status == "imported" {
|
||||
label = "旧配置迁入数字人(未重新蒸馏)"
|
||||
}
|
||||
fmt.Fprintf(&b, "\n\n%s\n## %s:%s(v%d)\n", start, label, human.Name, version.Number)
|
||||
b.WriteString("以下档案仅约束表达方式,不能替代任务目标、事实来源、工具权限或上文衔接约束。不得把档案样本当作新的事实依据。\n")
|
||||
for _, key := range keys {
|
||||
if value := strings.TrimSpace(version.Persona[key]); value != "" {
|
||||
fmt.Fprintf(&b, "%s:%s\n", key, value)
|
||||
}
|
||||
}
|
||||
b.WriteString(end)
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func syncDigitalHumanBindings(tx *gorm.DB, scope types.DigitalHumanScope, human *types.DigitalHuman, version *types.DigitalHumanVersion) error {
|
||||
var bindings []types.DigitalHumanBinding
|
||||
if err := tx.Where("tenant_id = ? AND digital_human_id = ?", scope.TenantID, human.ID).Find(&bindings).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, binding := range bindings {
|
||||
var agent types.CustomAgent
|
||||
if err := tx.First(&agent, "id = ? AND tenant_id = ? AND created_by = ?", binding.AgentID, scope.TenantID, binding.OwnerID).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
continue
|
||||
}
|
||||
return err
|
||||
}
|
||||
agent.Config.SystemPrompt = digitalHumanPrompt(human, version, agent.Config.SystemPrompt)
|
||||
if err := tx.Save(&agent).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
binding.AppliedVersion = version.Number
|
||||
if err := tx.Save(&binding).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Review(ctx context.Context, scope types.DigitalHumanScope, id string, number int, approve bool, note string) error {
|
||||
return r.mutate(ctx, scope, id, func(tx *gorm.DB, human *types.DigitalHuman) error {
|
||||
if human.ProcessingJobID != "" || human.LatestVersion != number {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
var version types.DigitalHumanVersion
|
||||
if err := dhScope(tx, scope).First(&version, "digital_human_id = ? AND number = ?", id, number).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
if version.Status != "draft" {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
version.Status = "rejected"
|
||||
if approve {
|
||||
hasContent := false
|
||||
for _, value := range version.Persona {
|
||||
hasContent = hasContent || strings.TrimSpace(value) != ""
|
||||
}
|
||||
if !hasContent {
|
||||
return &types.DigitalHumanInputError{Message: "人格档案为空,请先蒸馏或编辑"}
|
||||
}
|
||||
version.Status, human.Status, human.ActiveVersion = "approved", "active", number
|
||||
if err := syncDigitalHumanBindings(tx, scope, human, &version); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if human.ActiveVersion > 0 {
|
||||
human.Status = "active"
|
||||
} else {
|
||||
human.Status = "draft"
|
||||
}
|
||||
now := time.Now()
|
||||
version.ReviewedAt, version.ReviewNote = &now, note
|
||||
if err := tx.Save(&version).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
human.Revision++
|
||||
return tx.Save(human).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Bind(ctx context.Context, scope types.DigitalHumanScope, id, agentID string) error {
|
||||
return r.bindingTransaction(ctx, scope, id, agentID, false, func(tx *gorm.DB, human *types.DigitalHuman, agent *types.CustomAgent) error {
|
||||
var existing types.DigitalHumanBinding
|
||||
err := tx.Where("tenant_id = ? AND agent_id = ?", scope.TenantID, agentID).First(&existing).Error
|
||||
if err == nil {
|
||||
if existing.DigitalHumanID == id {
|
||||
return nil
|
||||
}
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return err
|
||||
}
|
||||
binding := types.DigitalHumanBinding{ID: uuid.NewString(), TenantID: scope.TenantID, OwnerID: agent.CreatedBy, DigitalHumanID: id, AgentID: agentID, BasePrompt: agent.Config.SystemPrompt}
|
||||
if human.ActiveVersion > 0 {
|
||||
var version types.DigitalHumanVersion
|
||||
if err := tx.Where("tenant_id = ? AND owner_id = ?", human.TenantID, human.OwnerID).First(&version, "digital_human_id = ? AND number = ?", id, human.ActiveVersion).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
agent.Config.SystemPrompt, binding.AppliedVersion = digitalHumanPrompt(human, &version, binding.BasePrompt), version.Number
|
||||
if err := tx.Save(agent).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Create(&binding).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Unbind(ctx context.Context, scope types.DigitalHumanScope, id, agentID string) error {
|
||||
return r.bindingTransaction(ctx, scope, id, agentID, true, func(tx *gorm.DB, human *types.DigitalHuman, agent *types.CustomAgent) error {
|
||||
if human.PrimaryAgentID == agentID {
|
||||
return &types.DigitalHumanInputError{Message: "默认同名 Agent 随数字人管理,请删除数字人时一并解绑"}
|
||||
}
|
||||
var binding types.DigitalHumanBinding
|
||||
if err := tx.Where("tenant_id = ? AND owner_id = ?", scope.TenantID, agent.CreatedBy).First(&binding, "digital_human_id = ? AND agent_id = ?", id, agentID).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
if err := restoreDigitalHumanAgent(tx, scope, &binding); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Delete(&binding).Error
|
||||
})
|
||||
}
|
||||
|
||||
// Mounting a persona changes the target Agent, not the shared persona. Lock
|
||||
// the human first so reviews/deletion cannot race a new binding.
|
||||
func (r *DigitalHumanRepository) bindingTransaction(ctx context.Context, scope types.DigitalHumanScope, id, agentID string, unbind bool, fn func(*gorm.DB, *types.DigitalHuman, *types.CustomAgent) error) error {
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var human types.DigitalHuman
|
||||
q := tx.Clauses(clause.Locking{Strength: "UPDATE"})
|
||||
if unbind {
|
||||
q = q.Where("tenant_id = ?", scope.TenantID)
|
||||
} else {
|
||||
q = dhReadScope(q, scope)
|
||||
}
|
||||
if err := q.First(&human, "id = ?", id).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
var agent types.CustomAgent
|
||||
q = tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id = ? AND id = ? AND is_builtin = ?", scope.TenantID, agentID, false)
|
||||
if !scope.ReadAll {
|
||||
q = q.Where("created_by = ?", scope.OwnerID)
|
||||
}
|
||||
if err := q.First(&agent).Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
return fn(tx, &human, &agent)
|
||||
})
|
||||
}
|
||||
|
||||
func restoreDigitalHumanAgent(tx *gorm.DB, scope types.DigitalHumanScope, binding *types.DigitalHumanBinding) error {
|
||||
var agent types.CustomAgent
|
||||
err := tx.First(&agent, "id = ? AND tenant_id = ? AND created_by = ?", binding.AgentID, scope.TenantID, binding.OwnerID).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if binding.AppliedVersion == 0 {
|
||||
return nil
|
||||
}
|
||||
agent.Config.SystemPrompt = removeDigitalHumanPrompt(agent.Config.SystemPrompt, binding.DigitalHumanID)
|
||||
return tx.Save(&agent).Error
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Delete(ctx context.Context, scope types.DigitalHumanScope, id string) error {
|
||||
_, err := r.DeleteWithIMChannels(ctx, scope, id)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) DeleteWithIMChannels(ctx context.Context, scope types.DigitalHumanScope, id string) ([]string, error) {
|
||||
var channelIDs []string
|
||||
err := r.mutate(ctx, scope, id, func(tx *gorm.DB, human *types.DigitalHuman) error {
|
||||
if human.ProcessingJobID != "" {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
var bindings []types.DigitalHumanBinding
|
||||
if err := tx.Where("tenant_id = ? AND digital_human_id = ?", scope.TenantID, id).Find(&bindings).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, binding := range bindings {
|
||||
if binding.AgentID != human.PrimaryAgentID {
|
||||
if err := restoreDigitalHumanAgent(tx, scope, &binding); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := tx.Where("id = ? AND tenant_id = ? AND created_by = ?", human.PrimaryAgentID, scope.TenantID, scope.OwnerID).Delete(&types.CustomAgent{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := detachFusionAgent(tx, human.PrimaryAgentID, scope.TenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
var err error
|
||||
channelIDs, err = softDeleteAgentIMChannels(tx, human.PrimaryAgentID, scope.TenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Where("tenant_id = ? AND digital_human_id = ?", scope.TenantID, id).Delete(&types.DigitalHumanBinding{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := dhScope(tx, scope).Where("digital_human_id = ?", id).Delete(&types.DigitalHumanVersion{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := dhScope(tx, scope).Where("digital_human_id = ?", id).Delete(&types.DigitalHumanJob{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Delete(human).Error
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return channelIDs, nil
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Sources(ctx context.Context, scope types.DigitalHumanScope, id string, number int) (*types.DigitalHumanVersion, error) {
|
||||
var version types.DigitalHumanVersion
|
||||
err := dhScope(r.db.WithContext(ctx), scope).First(&version, "digital_human_id = ? AND number = ?", id, number).Error
|
||||
return &version, dhNotFound(err)
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) StartJob(ctx context.Context, scope types.DigitalHumanScope, id string, revision int64, job *types.DigitalHumanJob) error {
|
||||
return r.mutate(ctx, scope, id, func(tx *gorm.DB, human *types.DigitalHuman) error {
|
||||
if human.ProcessingJobID != "" || human.Revision != revision {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
if err := tx.Create(job).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
human.ProcessingJobID, human.Status = job.ID, "distilling"
|
||||
human.Revision++
|
||||
return tx.Save(human).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) Job(ctx context.Context, scope types.DigitalHumanScope, id, jobID string) (*types.DigitalHumanJob, error) {
|
||||
var job types.DigitalHumanJob
|
||||
err := dhReadScope(r.db.WithContext(ctx), scope).First(&job, "id = ? AND digital_human_id = ?", jobID, id).Error
|
||||
return &job, dhNotFound(err)
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) JobProgress(ctx context.Context, scope types.DigitalHumanScope, jobID string, progress int, stage string) error {
|
||||
result := dhScope(r.db.WithContext(ctx).Model(&types.DigitalHumanJob{}), scope).
|
||||
Where("id = ? AND status = ? AND expires_at > ?", jobID, "running", time.Now()).
|
||||
Updates(map[string]interface{}{"progress": progress, "stage": stage, "expires_at": time.Now().Add(types.DigitalHumanJobLease)})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) DistillCheckpoint(ctx context.Context, scope types.DigitalHumanScope, id, fingerprint string) (*types.DigitalHumanDistillCheckpoint, error) {
|
||||
var job types.DigitalHumanJob
|
||||
err := dhScope(r.db.WithContext(ctx), scope).
|
||||
Where("digital_human_id = ? AND checkpoint_fingerprint = ? AND status = ? AND checkpoint IS NOT NULL", id, fingerprint, "failed").
|
||||
Order("updated_at DESC, id DESC").Take(&job).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return job.Checkpoint, err
|
||||
}
|
||||
|
||||
// A checkpoint is private retry state. Locking the human fences superseded
|
||||
// workers without turning partially distilled evidence into a visible version.
|
||||
func (r *DigitalHumanRepository) SaveDistillCheckpoint(ctx context.Context, scope types.DigitalHumanScope, id, jobID, fingerprint string, checkpoint *types.DigitalHumanDistillCheckpoint) error {
|
||||
if len(fingerprint) != 64 || checkpoint == nil || checkpoint.CompletedBatches < 0 || checkpoint.CompletedBatches > checkpoint.TotalBatches || checkpoint.TotalBatches < 1 {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
encoded, err := json.Marshal(checkpoint)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return r.mutate(ctx, scope, id, func(tx *gorm.DB, human *types.DigitalHuman) error {
|
||||
if human.ProcessingJobID != jobID {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
var job types.DigitalHumanJob
|
||||
if err := dhScope(tx, scope).Where("id = ? AND status = ? AND expires_at > ?", jobID, "running", time.Now()).Take(&job).Error; err != nil {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
if job.CheckpointFingerprint != "" && job.CheckpointFingerprint != fingerprint {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
if job.Checkpoint != nil && (job.Checkpoint.TotalBatches != checkpoint.TotalBatches || job.Checkpoint.CompletedBatches > checkpoint.CompletedBatches) {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
result := dhScope(tx.Model(&types.DigitalHumanJob{}), scope).
|
||||
Where("id = ? AND status = ? AND expires_at > ?", jobID, "running", time.Now()).Updates(map[string]interface{}{
|
||||
"checkpoint_fingerprint": fingerprint, "checkpoint": string(encoded), "expires_at": time.Now().Add(types.DigitalHumanJobLease),
|
||||
})
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
func (r *DigitalHumanRepository) FinishJob(ctx context.Context, scope types.DigitalHumanScope, id, jobID string, version *types.DigitalHumanVersion, failure string) error {
|
||||
return r.mutate(ctx, scope, id, func(tx *gorm.DB, human *types.DigitalHuman) error {
|
||||
if human.ProcessingJobID != jobID {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
var job types.DigitalHumanJob
|
||||
if err := dhScope(tx, scope).First(&job, "id = ? AND status = ?", jobID, "running").Error; err != nil {
|
||||
return dhNotFound(err)
|
||||
}
|
||||
job.Status, job.Stage, job.Error = "failed", "蒸馏失败", failure
|
||||
if version != nil && time.Now().After(job.ExpiresAt) {
|
||||
version = nil
|
||||
job.Error = "蒸馏已超时,请重新提交。已审核版本保持可用。"
|
||||
}
|
||||
if version != nil {
|
||||
version.ID, version.TenantID, version.OwnerID, version.DigitalHumanID = uuid.NewString(), scope.TenantID, scope.OwnerID, id
|
||||
version.Number, version.Status = human.LatestVersion+1, "draft"
|
||||
if err := tx.Create(version).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
human.LatestVersion, human.Status = version.Number, "testing"
|
||||
job.Status, job.Progress, job.Stage, job.Version = "completed", 100, "蒸馏完成,等待人工审核", version.Number
|
||||
} else if human.ActiveVersion > 0 {
|
||||
human.Status = "active"
|
||||
} else {
|
||||
human.Status = "draft"
|
||||
}
|
||||
if err := tx.Save(&job).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
human.ProcessingJobID = ""
|
||||
human.Revision++
|
||||
return tx.Save(human).Error
|
||||
})
|
||||
}
|
||||
|
||||
// Recovery is scoped to readable humans but finalization still uses each actual owner.
|
||||
func (r *DigitalHumanRepository) RecoverExpiredJobs(ctx context.Context, scope types.DigitalHumanScope) error {
|
||||
var jobs []types.DigitalHumanJob
|
||||
if err := dhReadScope(r.db.WithContext(ctx), scope).Where("status = ? AND expires_at <= ?", "running", time.Now()).Limit(100).Find(&jobs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, job := range jobs {
|
||||
ownerScope := types.DigitalHumanScope{TenantID: job.TenantID, OwnerID: job.OwnerID}
|
||||
err := r.FinishJob(ctx, ownerScope, job.DigitalHumanID, job.ID, nil, "蒸馏已超时或服务重启,请重新提交。已审核版本保持可用。")
|
||||
if err != nil && !errors.Is(err, types.ErrDigitalHumanConflict) && !errors.Is(err, types.ErrDigitalHumanNotFound) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,237 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||||
apperrors "github.com/Tencent/WeKnora/internal/errors"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// FilterDigitalHumanAgents protects persona-bearing native agents in ordinary
|
||||
// tenant access paths. Explicit organization shares use their verified share path.
|
||||
func (r *customAgentRepository) FilterDigitalHumanAgents(ctx context.Context, agents []*types.CustomAgent) ([]*types.CustomAgent, error) {
|
||||
if len(agents) == 0 {
|
||||
return agents, nil
|
||||
}
|
||||
ids := make([]string, 0, len(agents))
|
||||
for _, agent := range agents {
|
||||
ids = append(ids, agent.ID)
|
||||
}
|
||||
var bindings []types.DigitalHumanBinding
|
||||
if err := r.db.WithContext(ctx).Where("agent_id IN ?", ids).Find(&bindings).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
byAgent := make(map[string]types.DigitalHumanBinding, len(bindings))
|
||||
humanIDs := make([]string, 0, len(bindings))
|
||||
for _, binding := range bindings {
|
||||
byAgent[binding.AgentID] = binding
|
||||
humanIDs = append(humanIDs, binding.DigitalHumanID)
|
||||
}
|
||||
userID, _ := types.UserIDFromContext(ctx)
|
||||
tenantID, _ := types.TenantIDFromContext(ctx)
|
||||
var humans []types.DigitalHuman
|
||||
if len(humanIDs) > 0 {
|
||||
if err := r.db.WithContext(ctx).Select("id", "name", "active_version").Where("tenant_id = ? AND id IN ?", tenantID, humanIDs).Find(&humans).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
byHuman := make(map[string]types.DigitalHuman, len(humans))
|
||||
for _, human := range humans {
|
||||
byHuman[human.ID] = human
|
||||
}
|
||||
admin := types.IsSystemAdminFromContext(ctx) || types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin)
|
||||
result := make([]*types.CustomAgent, 0, len(agents))
|
||||
for _, agent := range agents {
|
||||
binding, bound := byAgent[agent.ID]
|
||||
// Copies preserve the persona snapshot but do not have a live binding.
|
||||
// Keep those snapshots private until the creator removes that persona.
|
||||
ownerID, ownerTenant := binding.OwnerID, binding.TenantID
|
||||
private := bound || strings.Contains(agent.Config.SystemPrompt, "<!-- weknora:digital-human:")
|
||||
if !bound {
|
||||
ownerID, ownerTenant = agent.CreatedBy, agent.TenantID
|
||||
}
|
||||
shared := accesspolicy.Readable(ctx, "agents", agent.ID, false, false) || accesspolicy.CanExecuteDependency(ctx, agent.TenantID, "agents", agent.ID, agent.CreatedBy) || (bound && accesspolicy.Readable(ctx, "digital-humans", binding.DigitalHumanID, false, false))
|
||||
if !private || (tenantID == ownerTenant && (admin || shared || (userID != "" && !types.IsSyntheticUserID(userID) && userID == ownerID))) {
|
||||
agent.DigitalHuman = nil
|
||||
if human, ok := byHuman[binding.DigitalHumanID]; bound && ok {
|
||||
agent.DigitalHuman = &types.AgentDigitalHumanSummary{ID: human.ID, Name: human.Name, AppliedVersion: binding.AppliedVersion, ActiveVersion: human.ActiveVersion}
|
||||
}
|
||||
result = append(result, agent)
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// Re-read the binding and persona before each execution. A user's installed
|
||||
// snapshot does not turn a revoked shared persona into their own resource.
|
||||
func (r *customAgentRepository) CheckDigitalHumanAgentExecution(ctx context.Context, agent *types.CustomAgent) error {
|
||||
if !strings.Contains(agent.Config.SystemPrompt, "<!-- weknora:digital-human:") {
|
||||
return nil
|
||||
}
|
||||
denied := apperrors.NewForbiddenError("此 Agent 的数字人已删除、解绑或授权已撤销,请重新选择可用数字人")
|
||||
if tenant, ok := types.TenantIDFromContext(ctx); !ok || tenant != agent.TenantID || tenant == 0 {
|
||||
return denied
|
||||
}
|
||||
var binding types.DigitalHumanBinding
|
||||
if err := r.db.WithContext(ctx).Where("tenant_id = ? AND agent_id = ? AND owner_id = ?", agent.TenantID, agent.ID, agent.CreatedBy).First(&binding).Error; err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return denied
|
||||
}
|
||||
return err
|
||||
}
|
||||
var human types.DigitalHuman
|
||||
if err := r.db.WithContext(ctx).Where("tenant_id = ? AND id = ?", agent.TenantID, binding.DigitalHumanID).First(&human).Error; err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return denied
|
||||
}
|
||||
return err
|
||||
}
|
||||
if !strings.Contains(agent.Config.SystemPrompt, fmt.Sprintf("<!-- weknora:digital-human:%s:start -->", human.ID)) || strings.Count(agent.Config.SystemPrompt, "<!-- weknora:digital-human:") != 2 {
|
||||
return denied
|
||||
}
|
||||
if binding.AppliedVersion <= 0 || binding.AppliedVersion != human.ActiveVersion || !accesspolicy.Allowed(ctx, "digital-humans.read") || accesspolicy.Hidden(ctx, "digital-humans", human.ID) {
|
||||
return denied
|
||||
}
|
||||
userID, _ := types.UserIDFromContext(ctx)
|
||||
admin := types.IsSystemAdminFromContext(ctx) || types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin)
|
||||
ownedDependency := human.OwnerID != "" && human.OwnerID == agent.CreatedBy
|
||||
if !admin && !ownedDependency && !accesspolicy.Readable(ctx, "digital-humans", human.ID, human.OwnerID == userID, false) {
|
||||
return denied
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *tenantSkillRepository) CheckDigitalHumanAgentExecution(ctx context.Context, agent *types.CustomAgent) error {
|
||||
return (&customAgentRepository{db: r.db}).CheckDigitalHumanAgentExecution(ctx, agent)
|
||||
}
|
||||
|
||||
// An allowed copy must keep its live persona binding so later review,
|
||||
// revocation and deletion still apply to the installed profile.
|
||||
func (r *customAgentRepository) CreateAgentCopy(ctx context.Context, sourceID string, target *types.CustomAgent) error {
|
||||
if !strings.Contains(target.Config.SystemPrompt, "<!-- weknora:digital-human:") {
|
||||
return r.CreateAgent(ctx, target)
|
||||
}
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var original types.DigitalHumanBinding
|
||||
if err := tx.Where("tenant_id = ? AND agent_id = ?", target.TenantID, sourceID).First(&original).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var human types.DigitalHuman
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id = ? AND id = ?", target.TenantID, original.DigitalHumanID).First(&human).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var source types.CustomAgent
|
||||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id = ? AND id = ?", target.TenantID, sourceID).First(&source).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
var current types.DigitalHumanBinding
|
||||
if err := tx.Where("tenant_id = ? AND agent_id = ?", target.TenantID, sourceID).First(¤t).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if current.DigitalHumanID != human.ID || current.AppliedVersion != human.ActiveVersion {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
repo := &customAgentRepository{db: tx}
|
||||
if err := repo.CheckDigitalHumanAgentCopy(ctx, sourceID, target.TenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := repo.CheckDigitalHumanAgentExecution(ctx, &source); err != nil {
|
||||
return err
|
||||
}
|
||||
if source.Config.SystemPrompt != target.Config.SystemPrompt {
|
||||
return types.ErrDigitalHumanConflict
|
||||
}
|
||||
if err := tx.Create(target).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
binding := types.DigitalHumanBinding{ID: uuid.NewString(), TenantID: target.TenantID, OwnerID: target.CreatedBy, DigitalHumanID: human.ID, AgentID: target.ID, BasePrompt: removeDigitalHumanPrompt(target.Config.SystemPrompt, human.ID), AppliedVersion: human.ActiveVersion}
|
||||
return tx.Create(&binding).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *customAgentRepository) CheckDigitalHumanAgentCopy(ctx context.Context, id string, tenantID uint64) error {
|
||||
if types.IsSystemAdminFromContext(ctx) || types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin) {
|
||||
return nil
|
||||
}
|
||||
userID, _ := types.UserIDFromContext(ctx)
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).Table("digital_human_bindings AS binding").
|
||||
Joins("JOIN digital_humans AS human ON human.id = binding.digital_human_id AND human.tenant_id = binding.tenant_id").
|
||||
Where("binding.agent_id = ? AND binding.tenant_id = ? AND human.owner_id <> ?", id, tenantID, userID).
|
||||
Count(&count).Error
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count > 0 {
|
||||
return apperrors.NewForbiddenError("此 Agent 挂载了共享数字人,只能使用,不能复制其档案")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *customAgentRepository) CheckDigitalHumanAgentDeletion(ctx context.Context, id string, tenantID uint64) error {
|
||||
var count int64
|
||||
if err := r.db.WithContext(ctx).Model(&types.DigitalHuman{}).Where("primary_agent_id = ? AND tenant_id = ?", id, tenantID).Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count > 0 {
|
||||
return types.ErrDigitalHumanPrimaryAgent
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *customAgentRepository) DeleteAgentWithDigitalHumanBindings(ctx context.Context, id string, tenantID uint64) ([]string, error) {
|
||||
var channelIDs []string
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
repo := &customAgentRepository{db: tx}
|
||||
if err := repo.CheckDigitalHumanAgentDeletion(ctx, id, tenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := repo.DeleteAgent(ctx, id, tenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := detachFusionAgent(tx, id, tenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
var err error
|
||||
channelIDs, err = softDeleteAgentIMChannels(tx, id, tenantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Where("agent_id = ? AND tenant_id = ?", id, tenantID).Delete(&types.DigitalHumanBinding{}).Error
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return channelIDs, nil
|
||||
}
|
||||
|
||||
func detachFusionAgent(tx *gorm.DB, id string, tenantID uint64) error {
|
||||
if tx.Dialector.Name() == "postgres" {
|
||||
return tx.Exec(`UPDATE custom_agents SET config = jsonb_set(config::jsonb, '{fusion_agent_ids}', (config::jsonb->'fusion_agent_ids') - ?), updated_at = CURRENT_TIMESTAMP WHERE tenant_id = ? AND deleted_at IS NULL AND jsonb_typeof(config::jsonb->'fusion_agent_ids') = 'array' AND jsonb_exists(config::jsonb->'fusion_agent_ids', ?)`, id, tenantID, id).Error
|
||||
}
|
||||
var agents []types.CustomAgent
|
||||
if err := tx.Where("tenant_id = ?", tenantID).Find(&agents).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, agent := range agents {
|
||||
ids := make([]string, 0, len(agent.Config.FusionAgentIDs))
|
||||
for _, fusionID := range agent.Config.FusionAgentIDs {
|
||||
if fusionID != id {
|
||||
ids = append(ids, fusionID)
|
||||
}
|
||||
}
|
||||
if len(ids) != len(agent.Config.FusionAgentIDs) {
|
||||
agent.Config.FusionAgentIDs = ids
|
||||
if err := tx.Model(&agent).Update("config", agent.Config).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,343 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupDigitalHumanTest(t *testing.T) (*DigitalHumanRepository, *gorm.DB, types.DigitalHumanScope, *types.DigitalHuman) {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { sqlDB.Close() })
|
||||
require.NoError(t, db.AutoMigrate(&types.CustomAgent{}, &types.DigitalHuman{}, &types.DigitalHumanVersion{}, &types.DigitalHumanBinding{}, &types.DigitalHumanJob{}))
|
||||
require.NoError(t, db.Exec("CREATE TABLE users (id TEXT PRIMARY KEY, username TEXT, deleted_at DATETIME)").Error)
|
||||
require.NoError(t, db.Exec("CREATE TABLE im_channels (id TEXT PRIMARY KEY, agent_id TEXT, tenant_id INTEGER, enabled BOOLEAN, updated_at DATETIME, deleted_at DATETIME)").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO users (id, username) VALUES (?, ?)", "owner", "Teacher Owner").Error)
|
||||
scope := types.DigitalHumanScope{TenantID: 7, OwnerID: "owner"}
|
||||
human := &types.DigitalHuman{ID: uuid.NewString(), TenantID: scope.TenantID, OwnerID: scope.OwnerID, Name: "Teacher", PrimaryAgentID: uuid.NewString(), LatestVersion: 1, Status: "draft", Revision: 1}
|
||||
repo := NewDigitalHumanRepository(db)
|
||||
require.NoError(t, repo.Create(context.Background(), human, map[string]string{"tone": "measured"}))
|
||||
return repo, db, scope, human
|
||||
}
|
||||
|
||||
func TestDigitalHumanProgressRenewsLeaseAndFencesExpiredOrForeignJobs(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.Background()
|
||||
job := &types.DigitalHumanJob{ID: "job", TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
|
||||
require.NoError(t, repo.StartJob(ctx, scope, human.ID, human.Revision, job))
|
||||
require.NoError(t, repo.JobProgress(ctx, scope, job.ID, 40, "batch 4/10"))
|
||||
stored, err := repo.Job(ctx, scope, human.ID, job.ID)
|
||||
require.NoError(t, err)
|
||||
require.Greater(t, stored.ExpiresAt.Sub(time.Now()), 9*time.Minute)
|
||||
foreign := scope
|
||||
foreign.OwnerID = "other-owner"
|
||||
require.ErrorIs(t, repo.JobProgress(ctx, foreign, job.ID, 50, "foreign"), types.ErrDigitalHumanConflict)
|
||||
require.NoError(t, db.Model(&types.DigitalHumanJob{}).Where("id = ?", job.ID).Update("expires_at", time.Now().Add(-time.Minute)).Error)
|
||||
require.ErrorIs(t, repo.JobProgress(ctx, scope, job.ID, 60, "late"), types.ErrDigitalHumanConflict)
|
||||
}
|
||||
|
||||
func TestDigitalHumanCheckpointIsPrivateOwnerScopedAndFenced(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.Background()
|
||||
key := strings.Repeat("a", 64)
|
||||
checkpoint := &types.DigitalHumanDistillCheckpoint{CompletedBatches: 1, TotalBatches: 3, Persona: map[string]string{"tone": "private draft"}}
|
||||
job := &types.DigitalHumanJob{ID: "checkpoint-job", TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
|
||||
require.NoError(t, repo.StartJob(ctx, scope, human.ID, human.Revision, job))
|
||||
require.NoError(t, repo.SaveDistillCheckpoint(ctx, scope, human.ID, job.ID, key, checkpoint))
|
||||
foreign := scope
|
||||
foreign.OwnerID = "other"
|
||||
require.Error(t, repo.SaveDistillCheckpoint(ctx, foreign, human.ID, job.ID, key, checkpoint))
|
||||
require.ErrorIs(t, repo.SaveDistillCheckpoint(ctx, scope, human.ID, job.ID, strings.Repeat("b", 64), checkpoint), types.ErrDigitalHumanConflict)
|
||||
require.ErrorIs(t, repo.SaveDistillCheckpoint(ctx, scope, human.ID, job.ID, key, &types.DigitalHumanDistillCheckpoint{CompletedBatches: 0, TotalBatches: 3}), types.ErrDigitalHumanConflict)
|
||||
require.NoError(t, db.Model(&types.DigitalHumanJob{}).Where("id = ?", job.ID).Update("expires_at", time.Now().Add(-time.Minute)).Error)
|
||||
require.ErrorIs(t, repo.SaveDistillCheckpoint(ctx, scope, human.ID, job.ID, key, checkpoint), types.ErrDigitalHumanConflict)
|
||||
require.NoError(t, repo.RecoverExpiredJobs(ctx, scope))
|
||||
found, err := repo.DistillCheckpoint(ctx, scope, human.ID, key)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, checkpoint, found)
|
||||
for _, lookup := range []struct {
|
||||
scope types.DigitalHumanScope
|
||||
human, fingerprint string
|
||||
}{
|
||||
{foreign, human.ID, key},
|
||||
{types.DigitalHumanScope{TenantID: 8, OwnerID: scope.OwnerID}, human.ID, key},
|
||||
{scope, "another-human", key},
|
||||
{scope, human.ID, strings.Repeat("c", 64)},
|
||||
} {
|
||||
found, err := repo.DistillCheckpoint(ctx, lookup.scope, lookup.human, lookup.fingerprint)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, found)
|
||||
}
|
||||
detail, err := repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
newJob := &types.DigitalHumanJob{ID: "new-job", TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
|
||||
require.NoError(t, repo.StartJob(ctx, scope, human.ID, detail.Revision, newJob))
|
||||
// Even a stale worker with a stale local lease may not change old retry state.
|
||||
require.ErrorIs(t, repo.SaveDistillCheckpoint(ctx, scope, human.ID, job.ID, key, checkpoint), types.ErrDigitalHumanConflict)
|
||||
require.NoError(t, repo.SaveDistillCheckpoint(ctx, scope, human.ID, newJob.ID, key, checkpoint))
|
||||
}
|
||||
|
||||
func TestDigitalHumanDeletionCommitsChannelsTogether(t *testing.T) {
|
||||
for _, rollback := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "commit", true: "rollback"}[rollback], func(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
for _, row := range []struct {
|
||||
id, agent string
|
||||
tenant uint64
|
||||
}{
|
||||
{"owned", human.PrimaryAgentID, scope.TenantID},
|
||||
{"other-agent", "other-agent", scope.TenantID},
|
||||
{"other-tenant", human.PrimaryAgentID, scope.TenantID + 1},
|
||||
} {
|
||||
require.NoError(t, db.Exec("INSERT INTO im_channels (id, agent_id, tenant_id, enabled) VALUES (?, ?, ?, true)", row.id, row.agent, row.tenant).Error)
|
||||
}
|
||||
if rollback {
|
||||
require.NoError(t, db.Exec("CREATE TRIGGER reject_digital_human_delete BEFORE UPDATE OF deleted_at ON digital_humans BEGIN SELECT RAISE(ABORT, 'simulated write failure'); END").Error)
|
||||
}
|
||||
ids, err := repo.DeleteWithIMChannels(context.Background(), scope, human.ID)
|
||||
var channels int64
|
||||
require.NoError(t, db.Table("im_channels").Where("deleted_at IS NULL AND enabled = true").Count(&channels).Error)
|
||||
var agents int64
|
||||
require.NoError(t, db.Model(&types.CustomAgent{}).Where("id = ?", human.PrimaryAgentID).Count(&agents).Error)
|
||||
if rollback {
|
||||
require.Error(t, err)
|
||||
require.Empty(t, ids)
|
||||
require.Equal(t, int64(3), channels)
|
||||
require.Equal(t, int64(1), agents)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"owned"}, ids)
|
||||
require.Equal(t, int64(2), channels)
|
||||
require.Zero(t, agents)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDigitalHumanOwnerIsolationAndAdminRead(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.Background()
|
||||
var agent types.CustomAgent
|
||||
require.NoError(t, db.First(&agent, "id = ?", human.PrimaryAgentID).Error)
|
||||
require.Equal(t, human.Name, agent.Name)
|
||||
require.Equal(t, scope.OwnerID, agent.CreatedBy)
|
||||
other := types.DigitalHumanScope{TenantID: scope.TenantID, OwnerID: "other"}
|
||||
_, err := repo.Get(ctx, other, human.ID)
|
||||
require.ErrorIs(t, err, types.ErrDigitalHumanNotFound)
|
||||
other.ReadAll = true
|
||||
detail, err := repo.Get(ctx, other, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, detail.Versions, 1)
|
||||
require.Len(t, detail.Bindings, 1)
|
||||
require.Equal(t, "Teacher Owner", detail.CreatorName)
|
||||
listed, err := repo.List(ctx, other, 0, 100)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, listed, 1)
|
||||
require.Equal(t, "Teacher Owner", listed[0].CreatorName)
|
||||
require.ErrorIs(t, repo.Delete(ctx, other, human.ID), types.ErrDigitalHumanNotFound)
|
||||
require.ErrorIs(t, repo.Review(ctx, other, human.ID, 1, true, ""), types.ErrDigitalHumanNotFound)
|
||||
other.TenantID++
|
||||
_, err = repo.Get(ctx, other, human.ID)
|
||||
require.ErrorIs(t, err, types.ErrDigitalHumanNotFound)
|
||||
}
|
||||
|
||||
func TestDigitalHumanReviewBindingsAndDeletionPreserveAgentEdits(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.Background()
|
||||
extra := types.CustomAgent{ID: uuid.NewString(), TenantID: scope.TenantID, CreatedBy: scope.OwnerID, Name: "Content", Config: types.CustomAgentConfig{SystemPrompt: "Original task rules"}}
|
||||
require.NoError(t, db.Create(&extra).Error)
|
||||
require.NoError(t, repo.Bind(ctx, scope, human.ID, extra.ID))
|
||||
require.NoError(t, repo.Review(ctx, scope, human.ID, 1, true, "approved"))
|
||||
require.NoError(t, db.First(&extra, "id = ?", extra.ID).Error)
|
||||
require.Contains(t, extra.Config.SystemPrompt, "measured")
|
||||
extra.Config.SystemPrompt = "Revised task rules\n" + extra.Config.SystemPrompt + "\nKeep continuity"
|
||||
require.NoError(t, db.Save(&extra).Error)
|
||||
h, err := repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, repo.Update(ctx, scope, human.ID, h.Revision, h.Name, h.Description, h.Avatar, h.ModelID, map[string]string{"tone": "clear"}))
|
||||
require.NoError(t, repo.Review(ctx, scope, human.ID, 2, true, ""))
|
||||
require.NoError(t, db.First(&extra, "id = ?", extra.ID).Error)
|
||||
require.Contains(t, extra.Config.SystemPrompt, "Revised task rules")
|
||||
require.Contains(t, extra.Config.SystemPrompt, "Keep continuity")
|
||||
require.NotContains(t, extra.Config.SystemPrompt, "measured")
|
||||
require.Equal(t, 1, strings.Count(extra.Config.SystemPrompt, ":start -->"))
|
||||
require.NoError(t, repo.Delete(ctx, scope, human.ID))
|
||||
require.NoError(t, db.First(&extra, "id = ?", extra.ID).Error)
|
||||
require.NotContains(t, extra.Config.SystemPrompt, "weknora:digital-human:")
|
||||
require.Contains(t, extra.Config.SystemPrompt, "Keep continuity")
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.CustomAgent{}).Where("id = ?", human.PrimaryAgentID).Count(&count).Error)
|
||||
require.Zero(t, count)
|
||||
require.NoError(t, db.Model(&types.DigitalHumanBinding{}).Where("digital_human_id = ?", human.ID).Count(&count).Error)
|
||||
require.Zero(t, count)
|
||||
}
|
||||
|
||||
func TestDigitalHumanAgentCopyKeepsLiveBinding(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, scope.TenantID)
|
||||
ctx = context.WithValue(ctx, types.UserIDContextKey, scope.OwnerID)
|
||||
require.NoError(t, repo.Review(ctx, scope, human.ID, 1, true, "approved"))
|
||||
var source types.CustomAgent
|
||||
require.NoError(t, db.First(&source, "id = ?", human.PrimaryAgentID).Error)
|
||||
copy := source
|
||||
copy.ID, copy.Name = uuid.NewString(), "Copied persona"
|
||||
agents := &customAgentRepository{db: db}
|
||||
require.NoError(t, agents.CreateAgentCopy(ctx, source.ID, ©))
|
||||
require.NoError(t, agents.CheckDigitalHumanAgentExecution(ctx, ©))
|
||||
detail, err := repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, detail.Bindings, 2)
|
||||
require.NoError(t, repo.Update(ctx, scope, human.ID, detail.Revision, detail.Name, detail.Description, detail.Avatar, detail.ModelID, map[string]string{"tone": "new voice"}))
|
||||
require.NoError(t, repo.Review(ctx, scope, human.ID, 2, true, "approved"))
|
||||
require.NoError(t, db.First(©, "id = ?", copy.ID).Error)
|
||||
require.Contains(t, copy.Config.SystemPrompt, "new voice")
|
||||
require.NoError(t, agents.CheckDigitalHumanAgentExecution(ctx, ©))
|
||||
staleCopy := source
|
||||
staleCopy.ID = uuid.NewString()
|
||||
require.ErrorIs(t, agents.CreateAgentCopy(ctx, source.ID, &staleCopy), types.ErrDigitalHumanConflict)
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.CustomAgent{}).Where("id = ?", staleCopy.ID).Count(&count).Error)
|
||||
require.Zero(t, count)
|
||||
require.NoError(t, repo.Unbind(ctx, scope, human.ID, copy.ID))
|
||||
require.Error(t, agents.CheckDigitalHumanAgentExecution(ctx, ©), "saved profile must stop executing after unbind")
|
||||
}
|
||||
|
||||
func TestSharedDigitalHumanCanBeMountedWithoutChangingSourceOrOtherUsers(t *testing.T) {
|
||||
repo, db, owner, human := setupDigitalHumanTest(t)
|
||||
ctx := accesspolicy.WithSnapshot(context.Background(), &accesspolicy.Snapshot{Publications: map[string]bool{"digital-humans:" + human.ID: true}})
|
||||
consumer := types.DigitalHumanScope{TenantID: owner.TenantID, OwnerID: "consumer"}
|
||||
for _, agent := range []types.CustomAgent{
|
||||
{ID: "consumer-agent", TenantID: owner.TenantID, CreatedBy: consumer.OwnerID, Name: "Consumer", Config: types.CustomAgentConfig{SystemPrompt: "Consumer task"}},
|
||||
{ID: "other-agent", TenantID: owner.TenantID, CreatedBy: "other", Name: "Other"},
|
||||
{ID: "foreign-agent", TenantID: owner.TenantID + 1, CreatedBy: consumer.OwnerID, Name: "Foreign"},
|
||||
} {
|
||||
require.NoError(t, db.Create(&agent).Error)
|
||||
}
|
||||
require.NoError(t, repo.Review(ctx, owner, human.ID, 1, true, ""))
|
||||
require.NoError(t, repo.Bind(ctx, consumer, human.ID, "consumer-agent"))
|
||||
require.ErrorIs(t, repo.Bind(ctx, consumer, human.ID, "other-agent"), types.ErrDigitalHumanNotFound)
|
||||
require.ErrorIs(t, repo.Bind(ctx, consumer, human.ID, "foreign-agent"), types.ErrDigitalHumanNotFound)
|
||||
require.ErrorIs(t, repo.Delete(ctx, consumer, human.ID), types.ErrDigitalHumanNotFound)
|
||||
detail, err := repo.Get(ctx, consumer, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, detail.Bindings, 1, "recipients only see bindings for their own Agents")
|
||||
require.Equal(t, "consumer-agent", detail.Bindings[0].AgentID)
|
||||
var agent types.CustomAgent
|
||||
require.NoError(t, db.First(&agent, "id = ?", "consumer-agent").Error)
|
||||
require.Contains(t, agent.Config.SystemPrompt, "measured")
|
||||
agentRepo := &customAgentRepository{db: db}
|
||||
consumerContext := context.WithValue(ctx, types.TenantIDContextKey, owner.TenantID)
|
||||
consumerContext = context.WithValue(consumerContext, types.UserIDContextKey, consumer.OwnerID)
|
||||
visible, err := agentRepo.FilterDigitalHumanAgents(consumerContext, []*types.CustomAgent{&agent})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, visible, 1)
|
||||
require.NotNil(t, visible[0].DigitalHuman)
|
||||
require.Equal(t, human.ID, visible[0].DigitalHuman.ID)
|
||||
require.Equal(t, human.Name, visible[0].DigitalHuman.Name)
|
||||
require.Equal(t, 1, visible[0].DigitalHuman.AppliedVersion)
|
||||
require.ErrorContains(t, agentRepo.CheckDigitalHumanAgentCopy(consumerContext, agent.ID, owner.TenantID), "共享数字人")
|
||||
require.NoError(t, agentRepo.CheckDigitalHumanAgentExecution(consumerContext, &agent))
|
||||
savedSnapshot := agent
|
||||
revokedExecution := accesspolicy.WithSnapshot(consumerContext, &accesspolicy.Snapshot{Publications: map[string]bool{"digital-humans:" + human.ID: false}})
|
||||
require.ErrorContains(t, agentRepo.CheckDigitalHumanAgentExecution(revokedExecution, &savedSnapshot), "授权已撤销")
|
||||
foreignExecution := context.WithValue(consumerContext, types.TenantIDContextKey, owner.TenantID+1)
|
||||
require.Error(t, agentRepo.CheckDigitalHumanAgentExecution(foreignExecution, &savedSnapshot))
|
||||
adminContext := context.WithValue(consumerContext, types.TenantRoleContextKey, types.TenantRoleAdmin)
|
||||
require.NoError(t, agentRepo.CheckDigitalHumanAgentCopy(adminContext, agent.ID, owner.TenantID))
|
||||
// Owner publication updates every installed binding, preserving its owner.
|
||||
require.NoError(t, repo.Update(ctx, owner, human.ID, detail.Revision, human.Name, "", "", "", map[string]string{"tone": "revised"}))
|
||||
require.NoError(t, repo.Review(ctx, owner, human.ID, 2, true, ""))
|
||||
require.NoError(t, db.First(&agent, "id = ?", "consumer-agent").Error)
|
||||
require.Contains(t, agent.Config.SystemPrompt, "revised")
|
||||
require.Equal(t, consumer.OwnerID, agent.CreatedBy)
|
||||
// Revocation must still permit removing a persona from one's own Agent.
|
||||
revoked := accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{})
|
||||
require.NoError(t, repo.Unbind(revoked, consumer, human.ID, agent.ID))
|
||||
require.ErrorContains(t, agentRepo.CheckDigitalHumanAgentExecution(consumerContext, &savedSnapshot), "解绑")
|
||||
require.NoError(t, db.First(&agent, "id = ?", "consumer-agent").Error)
|
||||
require.Equal(t, "Consumer task", agent.Config.SystemPrompt)
|
||||
require.ErrorIs(t, repo.Bind(revoked, consumer, human.ID, agent.ID), types.ErrDigitalHumanNotFound)
|
||||
require.NoError(t, repo.Bind(ctx, consumer, human.ID, agent.ID))
|
||||
require.NoError(t, repo.Delete(ctx, owner, human.ID))
|
||||
require.NoError(t, db.First(&agent, "id = ?", "consumer-agent").Error)
|
||||
require.Equal(t, "Consumer task", agent.Config.SystemPrompt)
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.DigitalHumanBinding{}).Where("digital_human_id = ?", human.ID).Count(&count).Error)
|
||||
require.Zero(t, count)
|
||||
}
|
||||
|
||||
func TestDigitalHumanJobRecoveryFencesOldWorker(t *testing.T) {
|
||||
repo, _, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.Background()
|
||||
job := &types.DigitalHumanJob{ID: uuid.NewString(), TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(-time.Minute)}
|
||||
require.NoError(t, repo.StartJob(ctx, scope, human.ID, human.Revision, job))
|
||||
require.ErrorIs(t, repo.StartJob(ctx, scope, human.ID, human.Revision, job), types.ErrDigitalHumanConflict)
|
||||
require.NoError(t, repo.RecoverExpiredJobs(ctx, scope))
|
||||
h, err := repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, h.ProcessingJobID)
|
||||
stored, err := repo.Job(ctx, scope, human.ID, job.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "failed", stored.Status)
|
||||
require.ErrorIs(t, repo.FinishJob(ctx, scope, human.ID, job.ID, &types.DigitalHumanVersion{Persona: map[string]string{"tone": "stale"}}, ""), types.ErrDigitalHumanConflict)
|
||||
job2 := &types.DigitalHumanJob{ID: uuid.NewString(), TenantID: scope.TenantID, OwnerID: scope.OwnerID, DigitalHumanID: human.ID, Status: "running", ExpiresAt: time.Now().Add(time.Minute)}
|
||||
require.NoError(t, repo.StartJob(ctx, scope, human.ID, h.Revision, job2))
|
||||
require.NoError(t, repo.FinishJob(ctx, scope, human.ID, job2.ID, &types.DigitalHumanVersion{Persona: map[string]string{"tone": "fresh"}}, ""))
|
||||
h, err = repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, h.LatestVersion)
|
||||
require.Equal(t, "fresh", h.Versions[0].Persona["tone"])
|
||||
}
|
||||
|
||||
func TestImportedDigitalHumanKeepsStyleUntilNewVersionReviewed(t *testing.T) {
|
||||
repo, db, scope, human := setupDigitalHumanTest(t)
|
||||
ctx := context.Background()
|
||||
persona := map[string]string{"persona_name": "Legacy teacher", "tone": "direct", "taboo": "no jargon", "vocab": "hello", "identity": "teacher", "interaction": "ask questions"}
|
||||
require.NoError(t, db.Model(&types.DigitalHumanVersion{}).Where("digital_human_id = ?", human.ID).
|
||||
Updates(&types.DigitalHumanVersion{Status: "imported", Persona: persona}).Error)
|
||||
require.NoError(t, db.Model(human).Updates(map[string]any{"status": "active", "active_version": 1}).Error)
|
||||
extra := types.CustomAgent{ID: uuid.NewString(), TenantID: scope.TenantID, CreatedBy: scope.OwnerID,
|
||||
Name: "Imported style consumer", Config: types.CustomAgentConfig{SystemPrompt: "Keep current task rules"}}
|
||||
require.NoError(t, db.Create(&extra).Error)
|
||||
require.NoError(t, repo.Bind(ctx, scope, human.ID, extra.ID))
|
||||
require.NoError(t, db.First(&extra, "id = ?", extra.ID).Error)
|
||||
require.Contains(t, extra.Config.SystemPrompt, "旧配置迁入数字人(未重新蒸馏)")
|
||||
require.NotContains(t, extra.Config.SystemPrompt, "已审核数字人")
|
||||
for _, value := range persona {
|
||||
require.Contains(t, extra.Config.SystemPrompt, value)
|
||||
}
|
||||
require.ErrorIs(t, repo.Review(ctx, scope, human.ID, 1, true, ""), types.ErrDigitalHumanConflict)
|
||||
detail, err := repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, detail.Versions[0].ReviewedAt)
|
||||
require.NoError(t, repo.Update(ctx, scope, human.ID, detail.Revision, human.Name, "", "", "", map[string]string{"tone": "new tone"}))
|
||||
require.NoError(t, db.First(&extra, "id = ?", extra.ID).Error)
|
||||
require.Contains(t, extra.Config.SystemPrompt, "tone:direct")
|
||||
require.NotContains(t, extra.Config.SystemPrompt, "new tone")
|
||||
require.NoError(t, repo.Review(ctx, scope, human.ID, 2, false, "revise"))
|
||||
detail, err = repo.Get(ctx, scope, human.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, detail.ActiveVersion)
|
||||
require.Equal(t, "active", detail.Status)
|
||||
require.NoError(t, repo.Update(ctx, scope, human.ID, detail.Revision, human.Name, "", "", "", map[string]string{"tone": "reviewed replacement"}))
|
||||
require.NoError(t, repo.Review(ctx, scope, human.ID, 3, true, ""))
|
||||
require.NoError(t, db.First(&extra, "id = ?", extra.ID).Error)
|
||||
require.Contains(t, extra.Config.SystemPrompt, "已审核数字人")
|
||||
require.Contains(t, extra.Config.SystemPrompt, "Keep current task rules")
|
||||
require.Contains(t, extra.Config.SystemPrompt, "reviewed replacement")
|
||||
require.NotContains(t, extra.Config.SystemPrompt, "tone:direct")
|
||||
require.Equal(t, 1, strings.Count(extra.Config.SystemPrompt, ":start -->"))
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type embedChannelRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewEmbedChannelRepository(db *gorm.DB) interfaces.EmbedChannelRepository {
|
||||
return &embedChannelRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) Create(ctx context.Context, ch *types.EmbedChannel) error {
|
||||
return r.db.WithContext(ctx).Create(ch).Error
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) GetByID(ctx context.Context, id string) (*types.EmbedChannel, error) {
|
||||
var ch types.EmbedChannel
|
||||
err := r.db.WithContext(ctx).Where("id = ?", id).First(&ch).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ch, nil
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) GetByPublishToken(ctx context.Context, token string) (*types.EmbedChannel, error) {
|
||||
var ch types.EmbedChannel
|
||||
err := r.db.WithContext(ctx).Where("publish_token = ?", token).First(&ch).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &ch, nil
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) ListByAgent(
|
||||
ctx context.Context, tenantID uint64, agentID string,
|
||||
) ([]*types.EmbedChannel, error) {
|
||||
var rows []*types.EmbedChannel
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND agent_id = ?", tenantID, agentID).
|
||||
Order("created_at DESC").
|
||||
Find(&rows).Error
|
||||
return rows, err
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) ListByTenant(
|
||||
ctx context.Context, tenantID uint64,
|
||||
) ([]*types.EmbedChannel, error) {
|
||||
var rows []*types.EmbedChannel
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ?", tenantID).
|
||||
Order("created_at DESC").
|
||||
Find(&rows).Error
|
||||
return rows, err
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) Update(ctx context.Context, ch *types.EmbedChannel) error {
|
||||
return r.db.WithContext(ctx).Save(ch).Error
|
||||
}
|
||||
|
||||
func (r *embedChannelRepository) Delete(ctx context.Context, tenantID uint64, id string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND id = ?", tenantID, id).
|
||||
Delete(&types.EmbedChannel{}).Error
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrKBShareNotFound = errors.New("knowledge base share not found")
|
||||
ErrKBShareAlreadyExists = errors.New("knowledge base already shared to this organization")
|
||||
)
|
||||
|
||||
// kbShareRepository implements KBShareRepository interface
|
||||
type kbShareRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewKBShareRepository creates a new knowledge base share repository
|
||||
func NewKBShareRepository(db *gorm.DB) interfaces.KBShareRepository {
|
||||
return &kbShareRepository{db: db}
|
||||
}
|
||||
|
||||
// Create creates a new share record
|
||||
func (r *kbShareRepository) Create(ctx context.Context, share *types.KnowledgeBaseShare) error {
|
||||
// Check if share already exists
|
||||
var count int64
|
||||
r.db.WithContext(ctx).Model(&types.KnowledgeBaseShare{}).
|
||||
Where("knowledge_base_id = ? AND organization_id = ? AND deleted_at IS NULL", share.KnowledgeBaseID, share.OrganizationID).
|
||||
Count(&count)
|
||||
|
||||
if count > 0 {
|
||||
return ErrKBShareAlreadyExists
|
||||
}
|
||||
|
||||
return r.db.WithContext(ctx).Create(share).Error
|
||||
}
|
||||
|
||||
// GetByID gets a share record by ID
|
||||
func (r *kbShareRepository) GetByID(ctx context.Context, id string) (*types.KnowledgeBaseShare, error) {
|
||||
var share types.KnowledgeBaseShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("id = ?", id).
|
||||
First(&share).Error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrKBShareNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &share, nil
|
||||
}
|
||||
|
||||
// GetByKBAndOrg gets a share record by knowledge base ID and organization ID
|
||||
func (r *kbShareRepository) GetByKBAndOrg(ctx context.Context, kbID string, orgID string) (*types.KnowledgeBaseShare, error) {
|
||||
var share types.KnowledgeBaseShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("knowledge_base_id = ? AND organization_id = ?", kbID, orgID).
|
||||
First(&share).Error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrKBShareNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &share, nil
|
||||
}
|
||||
|
||||
// Update updates a share record
|
||||
func (r *kbShareRepository) Update(ctx context.Context, share *types.KnowledgeBaseShare) error {
|
||||
return r.db.WithContext(ctx).Model(&types.KnowledgeBaseShare{}).
|
||||
Where("id = ?", share.ID).
|
||||
Updates(share).Error
|
||||
}
|
||||
|
||||
// Delete soft deletes a share record
|
||||
func (r *kbShareRepository) Delete(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).Where("id = ?", id).Delete(&types.KnowledgeBaseShare{}).Error
|
||||
}
|
||||
|
||||
// DeleteByKnowledgeBaseID soft deletes all share records for a knowledge base (e.g. when the KB is deleted)
|
||||
func (r *kbShareRepository) DeleteByKnowledgeBaseID(ctx context.Context, kbID string) error {
|
||||
return r.db.WithContext(ctx).Where("knowledge_base_id = ?", kbID).Delete(&types.KnowledgeBaseShare{}).Error
|
||||
}
|
||||
|
||||
// DeleteByOrganizationID soft deletes all share records for an organization (e.g. when the org is deleted)
|
||||
func (r *kbShareRepository) DeleteByOrganizationID(ctx context.Context, orgID string) error {
|
||||
return r.db.WithContext(ctx).Where("organization_id = ?", orgID).Delete(&types.KnowledgeBaseShare{}).Error
|
||||
}
|
||||
|
||||
// ListByKnowledgeBase lists all share records for a knowledge base
|
||||
func (r *kbShareRepository) ListByKnowledgeBase(ctx context.Context, kbID string) ([]*types.KnowledgeBaseShare, error) {
|
||||
var shares []*types.KnowledgeBaseShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Preload("Organization").
|
||||
Where("knowledge_base_id = ?", kbID).
|
||||
Order("created_at DESC").
|
||||
Find(&shares).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// ListByOrganization lists all share records for an organization.
|
||||
// Excludes shares whose knowledge base has been soft-deleted.
|
||||
func (r *kbShareRepository) ListByOrganization(ctx context.Context, orgID string) ([]*types.KnowledgeBaseShare, error) {
|
||||
var shares []*types.KnowledgeBaseShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN knowledge_bases ON knowledge_bases.id = kb_shares.knowledge_base_id AND knowledge_bases.deleted_at IS NULL").
|
||||
Preload("KnowledgeBase").
|
||||
Preload("Organization").
|
||||
Where("kb_shares.organization_id = ? AND kb_shares.deleted_at IS NULL", orgID).
|
||||
Order("kb_shares.created_at DESC").
|
||||
Find(&shares).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// ListByOrganizations lists all share records for the given organizations (batch).
|
||||
func (r *kbShareRepository) ListByOrganizations(ctx context.Context, orgIDs []string) ([]*types.KnowledgeBaseShare, error) {
|
||||
if len(orgIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var shares []*types.KnowledgeBaseShare
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN knowledge_bases ON knowledge_bases.id = kb_shares.knowledge_base_id AND knowledge_bases.deleted_at IS NULL").
|
||||
Preload("KnowledgeBase").
|
||||
Preload("Organization").
|
||||
Where("kb_shares.organization_id IN ? AND kb_shares.deleted_at IS NULL", orgIDs).
|
||||
Order("kb_shares.created_at DESC").
|
||||
Find(&shares).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// ListSharedKBsForTenant lists all knowledge bases shared to organizations that the tenant participates in.
|
||||
// Excludes shares for soft-deleted organizations and soft-deleted knowledge bases.
|
||||
func (r *kbShareRepository) ListSharedKBsForTenant(ctx context.Context, tenantID uint64) ([]*types.KnowledgeBaseShare, error) {
|
||||
var shares []*types.KnowledgeBaseShare
|
||||
|
||||
// Get shares for organizations the tenant is a member of; exclude deleted orgs and deleted KBs.
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN knowledge_bases ON knowledge_bases.id = kb_shares.knowledge_base_id AND knowledge_bases.deleted_at IS NULL").
|
||||
Preload("KnowledgeBase").
|
||||
Preload("Organization").
|
||||
Joins("JOIN organization_tenant_members otm ON otm.organization_id = kb_shares.organization_id").
|
||||
Joins("JOIN organizations ON organizations.id = kb_shares.organization_id AND organizations.deleted_at IS NULL").
|
||||
Where("otm.tenant_id = ?", tenantID).
|
||||
Where("kb_shares.deleted_at IS NULL").
|
||||
Order("kb_shares.created_at DESC").
|
||||
Find(&shares).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return shares, nil
|
||||
}
|
||||
|
||||
// CountSharesByKnowledgeBaseID counts the number of organizations a knowledge base is shared with
|
||||
func (r *kbShareRepository) CountSharesByKnowledgeBaseID(ctx context.Context, kbID string) (int64, error) {
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).Model(&types.KnowledgeBaseShare{}).
|
||||
Where("knowledge_base_id = ? AND deleted_at IS NULL", kbID).
|
||||
Count(&count).Error
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// CountSharesByKnowledgeBaseIDs counts the number of shares for multiple knowledge bases at once
|
||||
func (r *kbShareRepository) CountSharesByKnowledgeBaseIDs(ctx context.Context, kbIDs []string) (map[string]int64, error) {
|
||||
if len(kbIDs) == 0 {
|
||||
return make(map[string]int64), nil
|
||||
}
|
||||
|
||||
type result struct {
|
||||
KnowledgeBaseID string `gorm:"column:knowledge_base_id"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
}
|
||||
|
||||
var results []result
|
||||
err := r.db.WithContext(ctx).Model(&types.KnowledgeBaseShare{}).
|
||||
Select("knowledge_base_id, COUNT(*) as count").
|
||||
Where("knowledge_base_id IN ? AND deleted_at IS NULL", kbIDs).
|
||||
Group("knowledge_base_id").
|
||||
Find(&results).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
countMap := make(map[string]int64)
|
||||
for _, r := range results {
|
||||
countMap[r.KnowledgeBaseID] = r.Count
|
||||
}
|
||||
return countMap, nil
|
||||
}
|
||||
|
||||
// CountByOrganizations returns share counts per organization (only orgs in orgIDs). Excludes deleted KBs.
|
||||
func (r *kbShareRepository) CountByOrganizations(ctx context.Context, orgIDs []string) (map[string]int64, error) {
|
||||
if len(orgIDs) == 0 {
|
||||
return make(map[string]int64), nil
|
||||
}
|
||||
type row struct {
|
||||
OrgID string `gorm:"column:organization_id"`
|
||||
Count int64 `gorm:"column:count"`
|
||||
}
|
||||
var rows []row
|
||||
err := r.db.WithContext(ctx).Model(&types.KnowledgeBaseShare{}).
|
||||
Joins("JOIN knowledge_bases ON knowledge_bases.id = kb_shares.knowledge_base_id AND knowledge_bases.deleted_at IS NULL").
|
||||
Select("kb_shares.organization_id as organization_id, COUNT(*) as count").
|
||||
Where("kb_shares.organization_id IN ? AND kb_shares.deleted_at IS NULL", orgIDs).
|
||||
Group("kb_shares.organization_id").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]int64)
|
||||
for _, o := range orgIDs {
|
||||
out[o] = 0
|
||||
}
|
||||
for _, r := range rows {
|
||||
out[r.OrgID] = r.Count
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,40 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestCreateKnowledgeDefaultsCustomMetadataToEmptyObject(t *testing.T) {
|
||||
dsn := "file:" + uuid.New().String() + "?mode=memory&cache=shared"
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
require.NoError(t, db.AutoMigrate(&types.Knowledge{}))
|
||||
|
||||
repo := NewKnowledgeRepository(db)
|
||||
knowledge := &types.Knowledge{
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Type: "file",
|
||||
Title: "document.txt",
|
||||
ParseStatus: types.ParseStatusPending,
|
||||
EnableStatus: "disabled",
|
||||
}
|
||||
|
||||
require.NoError(t, repo.CreateKnowledge(context.Background(), knowledge))
|
||||
require.JSONEq(t, `{}`, string(knowledge.CustomMetadata))
|
||||
|
||||
persisted, err := repo.GetKnowledgeByID(context.Background(), knowledge.TenantID, knowledge.ID)
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{}`, string(persisted.CustomMetadata))
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestFindByDataSourceExternalID(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID uint64 = 50
|
||||
kbID := uuid.New().String()
|
||||
dsA := uuid.New().String()
|
||||
dsB := uuid.New().String()
|
||||
|
||||
insertRow := func(dsID, extID string) string {
|
||||
id := uuid.New().String()
|
||||
metadata := fmt.Sprintf(`{"datasource_id":%q,"external_id":%q}`, dsID, extID)
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges
|
||||
(id, tenant_id, knowledge_base_id, type, title, source, parse_status, metadata)
|
||||
VALUES (?, ?, ?, 'document', ?, 'feishu', 'completed', ?)
|
||||
`, id, tenantID, kbID, extID, metadata).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
idA := insertRow(dsA, "file:shared")
|
||||
_ = insertRow(dsB, "file:shared")
|
||||
|
||||
found, err := repo.FindByDataSourceExternalID(ctx, tenantID, kbID, dsA, "file:shared")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, found)
|
||||
assert.Equal(t, idA, found.ID)
|
||||
}
|
||||
|
||||
func TestFindByDataSourceExternalID_DeletedRowsExcluded(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID uint64 = 51
|
||||
kbID := uuid.New().String()
|
||||
dsID := uuid.New().String()
|
||||
|
||||
id := uuid.New().String()
|
||||
metadata := fmt.Sprintf(`{"datasource_id":%q,"external_id":%q}`, dsID, "file:gone")
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges
|
||||
(id, tenant_id, knowledge_base_id, type, title, source, parse_status, metadata, deleted_at)
|
||||
VALUES (?, ?, ?, 'document', 'gone', 'feishu', 'completed', ?, '2026-01-01 00:00:00')
|
||||
`, id, tenantID, kbID, metadata).Error)
|
||||
|
||||
found, err := repo.FindByDataSourceExternalID(ctx, tenantID, kbID, dsID, "file:gone")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, found)
|
||||
}
|
||||
|
||||
func TestHardDeleteKnowledge(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID uint64 = 52
|
||||
kbID := uuid.New().String()
|
||||
id := uuid.New().String()
|
||||
metadata := `{"datasource_id":"ds-1","external_id":"file:1"}`
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges
|
||||
(id, tenant_id, knowledge_base_id, type, title, source, parse_status, metadata, deleted_at)
|
||||
VALUES (?, ?, ?, 'document', 'file:1', 'feishu', 'completed', ?, '2026-01-01 00:00:00')
|
||||
`, id, tenantID, kbID, metadata).Error)
|
||||
|
||||
require.NoError(t, repo.HardDeleteKnowledge(ctx, tenantID, id))
|
||||
|
||||
var count int64
|
||||
require.NoError(t, db.Unscoped().Model(&struct{}{}).Table("knowledges").Where("id = ?", id).Count(&count).Error)
|
||||
assert.Equal(t, int64(0), count)
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCheckKnowledgeExists_FileHashIsScopedByFileType(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db)
|
||||
ctx := context.Background()
|
||||
tenantID := uint64(1)
|
||||
kbID := uuid.NewString()
|
||||
const fileHash = "same-content-hash"
|
||||
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, file_name, file_type, file_hash, parse_status)
|
||||
VALUES (?, ?, ?, 'file', 'document.md', 'document.md', 'md', ?, 'completed')
|
||||
`, uuid.NewString(), tenantID, kbID, fileHash).Error)
|
||||
|
||||
t.Run("same content with another file type is allowed", func(t *testing.T) {
|
||||
exists, knowledge, err := repo.CheckKnowledgeExists(ctx, tenantID, kbID, &types.KnowledgeCheckParams{
|
||||
Type: "file",
|
||||
FileHash: fileHash,
|
||||
FileType: "txt",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.False(t, exists)
|
||||
assert.Nil(t, knowledge)
|
||||
})
|
||||
|
||||
t.Run("same content and file type remains a duplicate", func(t *testing.T) {
|
||||
exists, knowledge, err := repo.CheckKnowledgeExists(ctx, tenantID, kbID, &types.KnowledgeCheckParams{
|
||||
Type: "file",
|
||||
FileHash: fileHash,
|
||||
FileType: "md",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.True(t, exists)
|
||||
require.NotNil(t, knowledge)
|
||||
assert.Equal(t, "md", knowledge.FileType)
|
||||
})
|
||||
|
||||
t.Run("file type matching is case-insensitive", func(t *testing.T) {
|
||||
exists, knowledge, err := repo.CheckKnowledgeExists(ctx, tenantID, kbID, &types.KnowledgeCheckParams{
|
||||
Type: "file",
|
||||
FileHash: fileHash,
|
||||
FileType: "MD",
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.True(t, exists)
|
||||
require.NotNil(t, knowledge)
|
||||
assert.Equal(t, "md", knowledge.FileType)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,381 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// knowledgesTestDDL mirrors the columns of `knowledges` that
|
||||
// SetFinalizing / FinalizeSubtask / UpdateKnowledge actually read or write.
|
||||
// We inline the DDL (instead of AutoMigrate) so the schema is explicit,
|
||||
// and we include pending_subtasks_count from migration 000056 plus the
|
||||
// processing/finalizing/completed columns the helpers care about.
|
||||
const knowledgesTestDDL = `
|
||||
CREATE TABLE IF NOT EXISTS knowledges (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
knowledge_base_id VARCHAR(36) NOT NULL,
|
||||
type VARCHAR(50) NOT NULL DEFAULT '',
|
||||
title VARCHAR(255) NOT NULL DEFAULT '',
|
||||
description TEXT,
|
||||
source VARCHAR(2048) NOT NULL DEFAULT '',
|
||||
parse_status VARCHAR(50) NOT NULL DEFAULT 'unprocessed',
|
||||
enable_status VARCHAR(50) NOT NULL DEFAULT 'enabled',
|
||||
embedding_model_id VARCHAR(64),
|
||||
file_name VARCHAR(255),
|
||||
folder_path VARCHAR(1024) NOT NULL DEFAULT '',
|
||||
file_type VARCHAR(50),
|
||||
file_size BIGINT,
|
||||
file_path TEXT,
|
||||
file_hash VARCHAR(64),
|
||||
storage_size BIGINT NOT NULL DEFAULT 0,
|
||||
metadata TEXT,
|
||||
tag_id VARCHAR(36),
|
||||
summary_status VARCHAR(32) DEFAULT 'none',
|
||||
last_faq_import_result TEXT DEFAULT NULL,
|
||||
channel VARCHAR(50) NOT NULL DEFAULT 'web',
|
||||
pending_subtasks_count INT NOT NULL DEFAULT 0,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
processed_at DATETIME,
|
||||
error_message TEXT,
|
||||
deleted_at DATETIME
|
||||
);
|
||||
`
|
||||
|
||||
// setupKnowledgeTestDB returns an in-memory SQLite db with the knowledges
|
||||
// table. SQLite has a single-writer constraint, so we cap MaxOpenConns at 1
|
||||
// and set a busy timeout: concurrent goroutines line up on the same
|
||||
// connection (just like production write workloads serialize at the row
|
||||
// level). This is enough to exercise the atomic semantics of the helpers
|
||||
// without flaking on "database table is locked".
|
||||
func setupKnowledgeTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
dsn := "file:" + uuid.New().String() + "?mode=memory&cache=shared&_busy_timeout=5000"
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
require.NoError(t, db.Exec(knowledgesTestDDL).Error)
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
// insertProcessingKnowledge seeds a row in `processing` state ready for a
|
||||
// SetFinalizing transition.
|
||||
func insertProcessingKnowledge(t *testing.T, db *gorm.DB) string {
|
||||
t.Helper()
|
||||
id := uuid.New().String()
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, source, parse_status, pending_subtasks_count)
|
||||
VALUES (?, 1, ?, 'document', 'finalize-test', 'manual', 'processing', 0)
|
||||
`, id, uuid.New().String()).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
// reloadKnowledgeRow returns the parse_status and pending_subtasks_count of
|
||||
// a row directly via raw SQL — bypasses any GORM hook noise.
|
||||
func reloadKnowledgeRow(t *testing.T, db *gorm.DB, id string) (status string, count int) {
|
||||
t.Helper()
|
||||
row := db.Raw(`SELECT parse_status, pending_subtasks_count FROM knowledges WHERE id = ?`, id).Row()
|
||||
require.NoError(t, row.Scan(&status, &count))
|
||||
return status, count
|
||||
}
|
||||
|
||||
func reloadKnowledgeErrorMessage(t *testing.T, db *gorm.DB, id string) string {
|
||||
t.Helper()
|
||||
var msg string
|
||||
require.NoError(t, db.Raw(`SELECT COALESCE(error_message, '') FROM knowledges WHERE id = ?`, id).Scan(&msg).Error)
|
||||
return msg
|
||||
}
|
||||
|
||||
func TestKnowledgeRepository_UpdateKnowledgeColumnsSanitizesErrorMessage(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db)
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
invalid := "parse failed " + string([]byte{0xef, 0xbc, 0x2e})
|
||||
|
||||
require.NoError(t, repo.UpdateKnowledgeColumns(context.Background(), id, map[string]interface{}{
|
||||
"error_message": invalid,
|
||||
}))
|
||||
|
||||
got := reloadKnowledgeErrorMessage(t, db, id)
|
||||
if !utf8.ValidString(got) {
|
||||
t.Fatalf("persisted error_message is invalid UTF-8: % x", []byte(got))
|
||||
}
|
||||
assert.Equal(t, "parse failed .", got)
|
||||
}
|
||||
|
||||
func insertKnowledgeWithStatus(t *testing.T, db *gorm.DB, status string, deleted bool) string {
|
||||
t.Helper()
|
||||
id := uuid.New().String()
|
||||
deletedAt := interface{}(nil)
|
||||
if deleted {
|
||||
deletedAt = "2026-06-16 12:00:00"
|
||||
}
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, source, parse_status, deleted_at)
|
||||
VALUES (?, 1, ?, 'document', 'delete-test', 'manual', ?, ?)
|
||||
`, id, uuid.New().String(), status, deletedAt).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
// TestFinalizeSubtask_Concurrent_ExactlyOnePromote spawns N goroutines that
|
||||
// each call FinalizeSubtask after SetFinalizing(N), and asserts:
|
||||
// - the counter ends at zero,
|
||||
// - parse_status is "completed",
|
||||
// - exactly one caller observed promoted=true.
|
||||
//
|
||||
// This is the behavior the original "stuck pending_subtasks_count" bug
|
||||
// violated: clobbered counters meant some callers saw a non-zero value
|
||||
// after the true count had reached zero, and none of them promoted.
|
||||
func TestFinalizeSubtask_Concurrent_ExactlyOnePromote(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const n = 20
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
|
||||
transitioned, err := repo.SetFinalizing(ctx, id, n)
|
||||
require.NoError(t, err)
|
||||
require.True(t, transitioned, "SetFinalizing should transition processing -> finalizing")
|
||||
|
||||
var promoteWins atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, promoted, ferr := repo.FinalizeSubtask(ctx, id)
|
||||
if ferr != nil {
|
||||
t.Errorf("FinalizeSubtask: %v", ferr)
|
||||
return
|
||||
}
|
||||
if promoted {
|
||||
promoteWins.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
assert.Equal(t, int32(1), promoteWins.Load(),
|
||||
"exactly one caller must observe promoted=true even under concurrent decrements")
|
||||
|
||||
status, count := reloadKnowledgeRow(t, db, id)
|
||||
assert.Equal(t, types.ParseStatusCompleted, status)
|
||||
assert.Equal(t, 0, count)
|
||||
}
|
||||
|
||||
// TestFinalizeSubtask_PartialDecrement_StaysFinalizing verifies the row
|
||||
// remains in "finalizing" with the expected residual count when fewer
|
||||
// callers decrement than were seeded — the promote guard must not fire
|
||||
// early.
|
||||
func TestFinalizeSubtask_PartialDecrement_StaysFinalizing(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
_, err := repo.SetFinalizing(ctx, id, 3)
|
||||
require.NoError(t, err)
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
_, promoted, ferr := repo.FinalizeSubtask(ctx, id)
|
||||
require.NoError(t, ferr)
|
||||
assert.False(t, promoted, "promote must not fire while count > 0")
|
||||
}
|
||||
|
||||
status, count := reloadKnowledgeRow(t, db, id)
|
||||
assert.Equal(t, types.ParseStatusFinalizing, status)
|
||||
assert.Equal(t, 1, count)
|
||||
}
|
||||
|
||||
// TestFinalizeSubtask_DecrementClampedAtZero verifies the safety-net
|
||||
// clamp on the decrement: extra calls past the seeded count must not
|
||||
// underflow pending_subtasks_count below zero. (Reconciliation's
|
||||
// shortfall-release loop relies on this.)
|
||||
func TestFinalizeSubtask_DecrementClampedAtZero(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
_, err := repo.SetFinalizing(ctx, id, 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
// First decrement drains the only slot and promotes.
|
||||
_, promoted, err := repo.FinalizeSubtask(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, promoted)
|
||||
|
||||
// Subsequent decrements must be no-ops, not underflow.
|
||||
for i := 0; i < 3; i++ {
|
||||
_, promoted, err := repo.FinalizeSubtask(ctx, id)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, promoted)
|
||||
}
|
||||
|
||||
status, count := reloadKnowledgeRow(t, db, id)
|
||||
assert.Equal(t, types.ParseStatusCompleted, status)
|
||||
assert.Equal(t, 0, count, "pending_subtasks_count must be clamped at zero")
|
||||
}
|
||||
|
||||
// TestSetFinalizingAndFinalizeSubtask_ClearStaleErrorMessage is the
|
||||
// regression test for stale error_message: a row that failed once keeps
|
||||
// error_message set, and both entering finalizing (a new attempt) and
|
||||
// promoting to completed (a successful finish) must clear it so the UI
|
||||
// no longer shows an outdated failure.
|
||||
func TestSetFinalizingAndFinalizeSubtask_ClearStaleErrorMessage(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
require.NoError(t, db.Exec(
|
||||
`UPDATE knowledges SET error_message = ? WHERE id = ?`,
|
||||
"Task interrupted due to application restart",
|
||||
id,
|
||||
).Error)
|
||||
|
||||
transitioned, err := repo.SetFinalizing(ctx, id, 1)
|
||||
require.NoError(t, err)
|
||||
require.True(t, transitioned)
|
||||
assert.Empty(t, reloadKnowledgeErrorMessage(t, db, id),
|
||||
"SetFinalizing must clear error_message from the previous attempt")
|
||||
|
||||
require.NoError(t, db.Exec(
|
||||
`UPDATE knowledges SET error_message = ? WHERE id = ?`,
|
||||
"stale finalizing failure",
|
||||
id,
|
||||
).Error)
|
||||
_, promoted, err := repo.FinalizeSubtask(ctx, id)
|
||||
require.NoError(t, err)
|
||||
require.True(t, promoted)
|
||||
assert.Empty(t, reloadKnowledgeErrorMessage(t, db, id),
|
||||
"promotion to completed must clear error_message")
|
||||
}
|
||||
|
||||
// TestUpdateKnowledge_DoesNotClobberPendingCounter is the regression test
|
||||
// for the original bug: a full-row Save with a stale in-memory counter
|
||||
// must not write that stale value back, otherwise it overwrites atomic
|
||||
// decrements made by other goroutines.
|
||||
//
|
||||
// Sequence:
|
||||
// 1. SetFinalizing(N=5) -> counter=5
|
||||
// 2. Caller A loads the row (sees counter=5)
|
||||
// 3. FinalizeSubtask runs concurrently and decrements to counter=4
|
||||
// 4. Caller A modifies an unrelated field (Title) and calls UpdateKnowledge
|
||||
// 5. Counter must still be 4 (not 5).
|
||||
func TestUpdateKnowledge_DoesNotClobberPendingCounter(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
_, err := repo.SetFinalizing(ctx, id, 5)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Step 2: caller A snapshots the row with counter=5 in memory.
|
||||
loaded, err := repo.GetKnowledgeByID(ctx, 1, id)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 5, loaded.PendingSubtasksCount)
|
||||
|
||||
// Step 3: an enrichment subtask decrements concurrently.
|
||||
_, _, err = repo.FinalizeSubtask(ctx, id)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Step 4: caller A persists an unrelated change. The in-memory copy
|
||||
// of PendingSubtasksCount is the STALE 5 — Save must NOT write it.
|
||||
loaded.Title = "renamed-after-stale-load"
|
||||
require.NoError(t, repo.UpdateKnowledge(ctx, loaded))
|
||||
|
||||
// Step 5: the live counter is still 4, not clobbered back to 5.
|
||||
status, count := reloadKnowledgeRow(t, db, id)
|
||||
assert.Equal(t, types.ParseStatusFinalizing, status)
|
||||
assert.Equal(t, 4, count,
|
||||
"UpdateKnowledge must omit pending_subtasks_count so a stale in-memory value cannot clobber atomic decrements")
|
||||
|
||||
// And the unrelated field WAS persisted.
|
||||
reloaded, err := repo.GetKnowledgeByID(ctx, 1, id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "renamed-after-stale-load", reloaded.Title)
|
||||
}
|
||||
|
||||
// TestUpdateKnowledge_PendingCounterOmittedOnReset verifies the inverse
|
||||
// case the reparse paths rely on: even setting PendingSubtasksCount=0
|
||||
// in memory and calling UpdateKnowledge does NOT persist that value.
|
||||
// Reparse must use UpdateKnowledgeColumn explicitly.
|
||||
func TestUpdateKnowledge_PendingCounterOmittedOnReset(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
id := insertProcessingKnowledge(t, db)
|
||||
_, err := repo.SetFinalizing(ctx, id, 7)
|
||||
require.NoError(t, err)
|
||||
|
||||
loaded, err := repo.GetKnowledgeByID(ctx, 1, id)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Caller tries to reset the counter via Save — this must be a no-op
|
||||
// for that column. The dedicated UpdateKnowledgeColumn is the only
|
||||
// path that actually writes pending_subtasks_count.
|
||||
loaded.PendingSubtasksCount = 0
|
||||
require.NoError(t, repo.UpdateKnowledge(ctx, loaded))
|
||||
|
||||
_, count := reloadKnowledgeRow(t, db, id)
|
||||
assert.Equal(t, 7, count, "UpdateKnowledge with PendingSubtasksCount=0 must NOT persist the reset")
|
||||
|
||||
// The explicit column write IS the supported path.
|
||||
require.NoError(t, repo.UpdateKnowledgeColumn(ctx, id, "pending_subtasks_count", 0))
|
||||
_, count = reloadKnowledgeRow(t, db, id)
|
||||
assert.Equal(t, 0, count)
|
||||
}
|
||||
|
||||
func TestUpdateActiveDeletingKnowledgeColumns_GuardsStateAndSoftDelete(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
activeDeletingID := insertKnowledgeWithStatus(t, db, types.ParseStatusDeleting, false)
|
||||
activeCompletedID := insertKnowledgeWithStatus(t, db, types.ParseStatusCompleted, false)
|
||||
deletedDeletingID := insertKnowledgeWithStatus(t, db, types.ParseStatusDeleting, true)
|
||||
|
||||
updated, err := repo.UpdateActiveDeletingKnowledgeColumns(ctx, activeDeletingID, map[string]interface{}{
|
||||
"parse_status": types.ParseStatusFailed,
|
||||
"error_message": "delete task exhausted retries",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.True(t, updated)
|
||||
|
||||
updated, err = repo.UpdateActiveDeletingKnowledgeColumns(ctx, activeCompletedID, map[string]interface{}{
|
||||
"parse_status": types.ParseStatusFailed,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, updated)
|
||||
|
||||
updated, err = repo.UpdateActiveDeletingKnowledgeColumns(ctx, deletedDeletingID, map[string]interface{}{
|
||||
"parse_status": types.ParseStatusFailed,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, updated)
|
||||
|
||||
status, _ := reloadKnowledgeRow(t, db, activeDeletingID)
|
||||
assert.Equal(t, types.ParseStatusFailed, status)
|
||||
status, _ = reloadKnowledgeRow(t, db, activeCompletedID)
|
||||
assert.Equal(t, types.ParseStatusCompleted, status)
|
||||
status, _ = reloadKnowledgeRow(t, db, deletedDeletingID)
|
||||
assert.Equal(t, types.ParseStatusDeleting, status)
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// folderPathOf reads a row's stored folder_path directly so the assertions do
|
||||
// not depend on the list filters.
|
||||
func folderPathOf(t *testing.T, db *gorm.DB, id string) string {
|
||||
t.Helper()
|
||||
var path string
|
||||
require.NoError(t, db.Raw(`SELECT folder_path FROM knowledges WHERE id = ?`, id).Scan(&path).Error)
|
||||
return path
|
||||
}
|
||||
|
||||
func TestUpdateKnowledgeFolderPath(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
otherKBID := uuid.New().String()
|
||||
|
||||
loose := insertKnowledgeInFolder(t, db, tenantID, kbID, "", "readme.md")
|
||||
filed := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs", "intro.md")
|
||||
foreign := insertKnowledgeInFolder(t, db, tenantID, otherKBID, "docs", "other.md")
|
||||
|
||||
affected, err := repo.UpdateKnowledgeFolderPath(ctx, tenantID, kbID, []string{loose, filed}, "archive/2026")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), affected)
|
||||
assert.Equal(t, "archive/2026", folderPathOf(t, db, loose))
|
||||
assert.Equal(t, "archive/2026", folderPathOf(t, db, filed))
|
||||
|
||||
// Moving back out of every folder is expressed as the empty path.
|
||||
affected, err = repo.UpdateKnowledgeFolderPath(ctx, tenantID, kbID, []string{filed}, "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), affected)
|
||||
assert.Equal(t, "", folderPathOf(t, db, filed))
|
||||
|
||||
// A row in another knowledge base is never touched, even when its ID is passed.
|
||||
affected, err = repo.UpdateKnowledgeFolderPath(ctx, tenantID, kbID, []string{foreign}, "leaked")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), affected)
|
||||
assert.Equal(t, "docs", folderPathOf(t, db, foreign))
|
||||
|
||||
affected, err = repo.UpdateKnowledgeFolderPath(ctx, tenantID, kbID, nil, "archive")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), affected)
|
||||
}
|
||||
|
||||
func TestRenameKnowledgeFolderPath(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
|
||||
root := insertKnowledgeInFolder(t, db, tenantID, kbID, "", "readme.md")
|
||||
docs := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs", "intro.md")
|
||||
spec := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs/spec", "design.md")
|
||||
deep := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs/spec/v2", "draft.md")
|
||||
sibling := insertKnowledgeInFolder(t, db, tenantID, kbID, "docsets", "unrelated.md")
|
||||
|
||||
// The folder and its whole subtree move; siblings whose name merely shares a
|
||||
// prefix ("docsets") must not be dragged along.
|
||||
affected, err := repo.RenameKnowledgeFolderPath(ctx, tenantID, kbID, "docs", "handbook")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(3), affected)
|
||||
assert.Equal(t, "handbook", folderPathOf(t, db, docs))
|
||||
assert.Equal(t, "handbook/spec", folderPathOf(t, db, spec))
|
||||
assert.Equal(t, "handbook/spec/v2", folderPathOf(t, db, deep))
|
||||
assert.Equal(t, "docsets", folderPathOf(t, db, sibling))
|
||||
assert.Equal(t, "", folderPathOf(t, db, root))
|
||||
|
||||
// Renaming a nested folder onto an existing one merges them.
|
||||
_, err = repo.RenameKnowledgeFolderPath(ctx, tenantID, kbID, "handbook/spec", "docsets")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "docsets", folderPathOf(t, db, spec))
|
||||
assert.Equal(t, "docsets/v2", folderPathOf(t, db, deep))
|
||||
assert.Equal(t, "docsets", folderPathOf(t, db, sibling))
|
||||
|
||||
// An unknown folder is a no-op rather than an error.
|
||||
affected, err = repo.RenameKnowledgeFolderPath(ctx, tenantID, kbID, "missing", "whatever")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), affected)
|
||||
|
||||
_, err = repo.RenameKnowledgeFolderPath(ctx, tenantID, kbID, "", "whatever")
|
||||
assert.Error(t, err, "renaming the root itself is meaningless and must be rejected")
|
||||
}
|
||||
|
||||
func TestRenameKnowledgeFolderPathUnderscoreInPath(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
|
||||
parent := insertKnowledgeInFolder(t, db, tenantID, kbID, "my_docs", "intro.md")
|
||||
child := insertKnowledgeInFolder(t, db, tenantID, kbID, "my_docs/spec", "design.md")
|
||||
|
||||
affected, err := repo.RenameKnowledgeFolderPath(ctx, tenantID, kbID, "my_docs", "handbook")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), affected)
|
||||
assert.Equal(t, "handbook", folderPathOf(t, db, parent))
|
||||
assert.Equal(t, "handbook/spec", folderPathOf(t, db, child))
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// insertKnowledgeInFolder seeds a completed knowledge row stored in the given
|
||||
// folder so the folder filter and folder aggregation can be asserted.
|
||||
func insertKnowledgeInFolder(t *testing.T, db *gorm.DB, tenantID uint64, kbID, folderPath, fileName string) string {
|
||||
t.Helper()
|
||||
id := uuid.New().String()
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, source, parse_status, file_name, folder_path)
|
||||
VALUES (?, ?, ?, 'file', ?, 'manual', 'completed', ?, ?)
|
||||
`, id, tenantID, kbID, fileName, fileName, folderPath).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
func TestListPagedKnowledgeFolderScope(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
|
||||
rootID := insertKnowledgeInFolder(t, db, tenantID, kbID, "", "readme.md")
|
||||
docsID := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs", "intro.md")
|
||||
specID := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs/spec", "design.md")
|
||||
otherID := insertKnowledgeInFolder(t, db, tenantID, kbID, "assets", "logo.png")
|
||||
|
||||
page := &types.Pagination{Page: 1, PageSize: 100}
|
||||
list := func(filter types.KnowledgeListFilter) ([]string, int64) {
|
||||
rows, total, err := repo.ListPagedKnowledgeByKnowledgeBaseID(ctx, tenantID, kbID, page, filter)
|
||||
require.NoError(t, err)
|
||||
ids := make([]string, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
ids = append(ids, row.ID)
|
||||
}
|
||||
return ids, total
|
||||
}
|
||||
|
||||
// No folder scope: the flat listing is unchanged, which keeps every existing
|
||||
// caller (and the "all documents" view) working.
|
||||
ids, total := list(types.KnowledgeListFilter{})
|
||||
assert.ElementsMatch(t, []string{rootID, docsID, specID, otherID}, ids)
|
||||
assert.Equal(t, int64(4), total)
|
||||
|
||||
// Exact scope with an empty path means "knowledge base root only".
|
||||
ids, total = list(types.KnowledgeListFilter{FolderScope: types.FolderScopeExact})
|
||||
assert.Equal(t, []string{rootID}, ids)
|
||||
assert.Equal(t, int64(1), total)
|
||||
|
||||
ids, total = list(types.KnowledgeListFilter{FolderPath: "docs", FolderScope: types.FolderScopeExact})
|
||||
assert.Equal(t, []string{docsID}, ids)
|
||||
assert.Equal(t, int64(1), total)
|
||||
|
||||
ids, total = list(types.KnowledgeListFilter{FolderPath: "docs", FolderScope: types.FolderScopeSubtree})
|
||||
assert.ElementsMatch(t, []string{docsID, specID}, ids)
|
||||
assert.Equal(t, int64(2), total)
|
||||
|
||||
// Subtree scope at the root spans the whole knowledge base.
|
||||
ids, _ = list(types.KnowledgeListFilter{FolderScope: types.FolderScopeSubtree})
|
||||
assert.ElementsMatch(t, []string{rootID, docsID, specID, otherID}, ids)
|
||||
|
||||
// Folder scope composes with the other filter dimensions.
|
||||
ids, _ = list(types.KnowledgeListFilter{
|
||||
FolderPath: "docs",
|
||||
FolderScope: types.FolderScopeSubtree,
|
||||
Keyword: "design",
|
||||
})
|
||||
assert.Equal(t, []string{specID}, ids)
|
||||
}
|
||||
|
||||
func TestListPagedKnowledgeFolderScopeUnderscoreInPath(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
|
||||
parentID := insertKnowledgeInFolder(t, db, tenantID, kbID, "my_docs", "intro.md")
|
||||
childID := insertKnowledgeInFolder(t, db, tenantID, kbID, "my_docs/spec", "design.md")
|
||||
// A sibling whose name merely shares a prefix must not be swept into subtree scope.
|
||||
lookalikeID := insertKnowledgeInFolder(t, db, tenantID, kbID, "myXdocs/spec", "other.md")
|
||||
|
||||
page := &types.Pagination{Page: 1, PageSize: 100}
|
||||
rows, total, err := repo.ListPagedKnowledgeByKnowledgeBaseID(ctx, tenantID, kbID, page, types.KnowledgeListFilter{
|
||||
FolderPath: "my_docs",
|
||||
FolderScope: types.FolderScopeSubtree,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
ids := make([]string, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
ids = append(ids, row.ID)
|
||||
}
|
||||
assert.ElementsMatch(t, []string{parentID, childID}, ids)
|
||||
assert.Equal(t, int64(2), total)
|
||||
assert.NotContains(t, ids, lookalikeID)
|
||||
}
|
||||
|
||||
func TestListKnowledgeFolderCounts(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
otherKBID := uuid.New().String()
|
||||
|
||||
insertKnowledgeInFolder(t, db, tenantID, kbID, "", "readme.md")
|
||||
insertKnowledgeInFolder(t, db, tenantID, kbID, "docs/spec", "design.md")
|
||||
insertKnowledgeInFolder(t, db, tenantID, kbID, "docs/spec", "api.md")
|
||||
insertKnowledgeInFolder(t, db, tenantID, otherKBID, "docs/spec", "other-kb.md")
|
||||
// Rows mid-deletion are excluded so the tree counts agree with the list.
|
||||
deleting := insertKnowledgeInFolder(t, db, tenantID, kbID, "docs", "vanishing.md")
|
||||
require.NoError(t, db.Exec(`UPDATE knowledges SET parse_status = ? WHERE id = ?`,
|
||||
types.ParseStatusDeleting, deleting).Error)
|
||||
|
||||
counts, err := repo.ListKnowledgeFolderCounts(ctx, tenantID, kbID)
|
||||
require.NoError(t, err)
|
||||
|
||||
byPath := map[string]int64{}
|
||||
for _, row := range counts {
|
||||
byPath[row.FolderPath] = row.Count
|
||||
}
|
||||
assert.Equal(t, map[string]int64{"": 1, "docs/spec": 2}, byPath)
|
||||
|
||||
tree := types.BuildKnowledgeFolderTree(counts)
|
||||
assert.Equal(t, int64(1), tree.RootDocumentCount)
|
||||
assert.Equal(t, int64(3), tree.TotalDocumentCount)
|
||||
require.Len(t, tree.Folders, 1)
|
||||
assert.Equal(t, "docs", tree.Folders[0].Path)
|
||||
assert.Equal(t, int64(2), tree.Folders[0].TotalCount)
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// insertKnowledgeInKB seeds a knowledge row in the given KB with the given
|
||||
// parse_status so we can assert the list filter's visibility rules.
|
||||
func insertKnowledgeInKB(t *testing.T, db *gorm.DB, tenantID uint64, kbID, status string) string {
|
||||
t.Helper()
|
||||
id := uuid.New().String()
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, source, parse_status)
|
||||
VALUES (?, ?, ?, 'file', 'list-filter-test', 'manual', ?)
|
||||
`, id, tenantID, kbID, status).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
// TestListPaged_ExcludesDeletingByDefault documents the fix for issue #2192:
|
||||
// a knowledge row that is mid-deletion (parse_status = 'deleting') must not
|
||||
// appear in the default document list, otherwise a successful async delete
|
||||
// keeps showing the entry until the row is physically gone. Rows in other
|
||||
// states (including the terminal 'failed' a stuck delete resolves to) stay
|
||||
// visible, and an explicit parse_status=deleting filter can still surface them.
|
||||
func TestListPaged_ExcludesDeletingByDefault(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID = uint64(1)
|
||||
kbID := uuid.New().String()
|
||||
|
||||
completedID := insertKnowledgeInKB(t, db, tenantID, kbID, "completed")
|
||||
failedID := insertKnowledgeInKB(t, db, tenantID, kbID, "failed")
|
||||
deletingID := insertKnowledgeInKB(t, db, tenantID, kbID, types.ParseStatusDeleting)
|
||||
|
||||
page := &types.Pagination{Page: 1, PageSize: 100}
|
||||
|
||||
// Default listing (no explicit parse_status): deleting rows are hidden.
|
||||
rows, total, err := repo.ListPagedKnowledgeByKnowledgeBaseID(
|
||||
ctx, tenantID, kbID, page, types.KnowledgeListFilter{},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
ids := make([]string, 0, len(rows))
|
||||
for _, r := range rows {
|
||||
ids = append(ids, r.ID)
|
||||
}
|
||||
assert.ElementsMatch(t, []string{completedID, failedID}, ids,
|
||||
"default list must include completed + failed but exclude deleting")
|
||||
assert.Equal(t, int64(2), total, "count must match the filtered rows")
|
||||
|
||||
// Explicit parse_status=deleting still surfaces the in-flight row so
|
||||
// operators / tooling can inspect stuck deletes.
|
||||
rows, total, err = repo.ListPagedKnowledgeByKnowledgeBaseID(
|
||||
ctx, tenantID, kbID, page,
|
||||
types.KnowledgeListFilter{ParseStatus: types.ParseStatusDeleting},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, rows, 1)
|
||||
assert.Equal(t, deletingID, rows[0].ID)
|
||||
assert.Equal(t, int64(1), total)
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestFindByMetadataKeyPrefix verifies that FindByMetadataKeyPrefix returns only
|
||||
// knowledge items whose metadata["external_id"] starts with the given prefix.
|
||||
//
|
||||
// Setup: three rows in the same KB:
|
||||
// - "nodeA" — the parent document; must NOT be matched by prefix "nodeA#"
|
||||
// - "nodeA#file#x" — an attachment child; MUST be matched
|
||||
// - "nodeB" — a sibling document; must NOT be matched
|
||||
//
|
||||
// A fourth row in a different KB with a matching external_id is also inserted to
|
||||
// verify tenant/KB isolation.
|
||||
func TestFindByMetadataKeyPrefix(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID uint64 = 42
|
||||
kbID := uuid.New().String()
|
||||
otherKBID := uuid.New().String()
|
||||
|
||||
insertRow := func(tid uint64, kb, extID string) string {
|
||||
id := uuid.New().String()
|
||||
metadata := fmt.Sprintf(`{"external_id":%q}`, extID)
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges
|
||||
(id, tenant_id, knowledge_base_id, type, title, source, parse_status, metadata)
|
||||
VALUES (?, ?, ?, 'document', ?, 'feishu', 'completed', ?)
|
||||
`, id, tid, kb, extID, metadata).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
_ = insertRow(tenantID, kbID, "nodeA") // parent — must NOT match
|
||||
childID := insertRow(tenantID, kbID, "nodeA#file#x") // attachment child — MUST match
|
||||
_ = insertRow(tenantID, kbID, "nodeB") // sibling — must NOT match
|
||||
_ = insertRow(tenantID, otherKBID, "nodeA#file#y") // different KB — must be excluded
|
||||
|
||||
results, err := repo.FindByMetadataKeyPrefix(ctx, tenantID, kbID, "external_id", "nodeA#")
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, results, 1, "only the attachment child should match the prefix 'nodeA#'")
|
||||
assert.Equal(t, childID, results[0].ID)
|
||||
}
|
||||
|
||||
// TestFindByMetadataKeyPrefix_DeletedRowsExcluded verifies that soft-deleted rows
|
||||
// (deleted_at IS NOT NULL) are excluded from the prefix query.
|
||||
func TestFindByMetadataKeyPrefix_DeletedRowsExcluded(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID uint64 = 43
|
||||
kbID := uuid.New().String()
|
||||
|
||||
id := uuid.New().String()
|
||||
metadata := fmt.Sprintf(`{"external_id":%q}`, "nodeC#file#z")
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges
|
||||
(id, tenant_id, knowledge_base_id, type, title, source, parse_status, metadata, deleted_at)
|
||||
VALUES (?, ?, ?, 'document', 'deleted-child', 'feishu', 'completed', ?, '2026-01-01 00:00:00')
|
||||
`, id, tenantID, kbID, metadata).Error)
|
||||
|
||||
results, err := repo.FindByMetadataKeyPrefix(ctx, tenantID, kbID, "external_id", "nodeC#")
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, results, "soft-deleted rows must not appear in prefix results")
|
||||
}
|
||||
|
||||
// TestFindByMetadataKeyPrefix_UnderscoreIsLiteral verifies that an underscore in
|
||||
// the prefix is treated as a literal character and does NOT act as a SQL LIKE
|
||||
// wildcard. Without escapeLikeKeyword applied to the prefix, a query for "ab_c#"
|
||||
// would also match "abXc#file#y" because '_' is a wildcard.
|
||||
func TestFindByMetadataKeyPrefix_UnderscoreIsLiteral(t *testing.T) {
|
||||
db := setupKnowledgeTestDB(t)
|
||||
repo := NewKnowledgeRepository(db).(*knowledgeRepository)
|
||||
ctx := context.Background()
|
||||
|
||||
const tenantID uint64 = 44
|
||||
kbID := uuid.New().String()
|
||||
|
||||
insertRow := func(extID string) string {
|
||||
id := uuid.New().String()
|
||||
metadata := fmt.Sprintf(`{"external_id":%q}`, extID)
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges
|
||||
(id, tenant_id, knowledge_base_id, type, title, source, parse_status, metadata)
|
||||
VALUES (?, ?, ?, 'document', ?, 'feishu', 'completed', ?)
|
||||
`, id, tenantID, kbID, extID, metadata).Error)
|
||||
return id
|
||||
}
|
||||
|
||||
// The prefix we will search for is "ab_c#". Without escaping, '_' is a
|
||||
// wildcard that matches any single character, so "abXc#file#y" would also
|
||||
// be returned. With escaping it must NOT be returned.
|
||||
targetID := insertRow("ab_c#file#x") // literal underscore — MUST match
|
||||
_ = insertRow("abXc#file#y") // 'X' instead of '_' — must NOT match
|
||||
|
||||
results, err := repo.FindByMetadataKeyPrefix(ctx, tenantID, kbID, "external_id", "ab_c#")
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Len(t, results, 1, "underscore in prefix must be treated as literal: only 'ab_c#file#x' should match")
|
||||
assert.Equal(t, targetID, results[0].ID)
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestKnowledgeSourceSchemaAllowsObjectStorageURLs(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:knowledge_source_schema?mode=memory&cache=shared"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
t.Fatalf("get sqlite handle: %v", err)
|
||||
}
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
if err := db.AutoMigrate(&types.Knowledge{}); err != nil {
|
||||
t.Fatalf("auto migrate knowledge: %v", err)
|
||||
}
|
||||
|
||||
var sourceType string
|
||||
rows, err := db.Raw("PRAGMA table_info(knowledges)").Rows()
|
||||
if err != nil {
|
||||
t.Fatalf("inspect knowledge schema: %v", err)
|
||||
}
|
||||
|
||||
for rows.Next() {
|
||||
var cid int
|
||||
var name, typ string
|
||||
var notNull int
|
||||
var defaultValue any
|
||||
var pk int
|
||||
if err := rows.Scan(&cid, &name, &typ, ¬Null, &defaultValue, &pk); err != nil {
|
||||
t.Fatalf("scan column info: %v", err)
|
||||
}
|
||||
if name == "source" {
|
||||
sourceType = strings.ToLower(typ)
|
||||
break
|
||||
}
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
t.Fatalf("close column info rows: %v", err)
|
||||
}
|
||||
if sourceType != "varchar(2048)" {
|
||||
t.Fatalf("knowledge source column type = %q, want varchar(2048)", sourceType)
|
||||
}
|
||||
|
||||
longURL := "https://example-bucket.cos.ap-beijing.myqcloud.com/test/" +
|
||||
strings.Repeat("encoded-path-segment-", 20) + ".docx"
|
||||
knowledge := &types.Knowledge{
|
||||
TenantID: 1,
|
||||
KnowledgeBaseID: "kb-1",
|
||||
Type: types.KnowledgeTypeManual,
|
||||
Title: "long source",
|
||||
Source: longURL,
|
||||
ParseStatus: types.ParseStatusPending,
|
||||
EnableStatus: "enabled",
|
||||
}
|
||||
if err := db.Create(knowledge).Error; err != nil {
|
||||
t.Fatalf("create knowledge with long source: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,271 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/common"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// KnowledgeSpanRepository persists the per-attempt span tree used by the
|
||||
// processing pipeline. Operations are deliberately narrow:
|
||||
//
|
||||
// - Upsert covers Begin/End/Fail/Skip — every state transition routes
|
||||
// through the same write so the row stays internally consistent.
|
||||
// - NextAttempt allocates a new attempt for re-parses without touching
|
||||
// historical rows. Old attempts stay queryable for post-mortem.
|
||||
// - ListByAttempt is the only read path; the handler builds the tree
|
||||
// in memory rather than recursing through the DB.
|
||||
type KnowledgeSpanRepository interface {
|
||||
Upsert(ctx context.Context, row *types.KnowledgeProcessingSpan) error
|
||||
NextAttempt(ctx context.Context, knowledgeID string) (int, error)
|
||||
LatestAttempt(ctx context.Context, knowledgeID string) (int, error)
|
||||
ListByAttempt(ctx context.Context, knowledgeID string, attempt int) ([]types.KnowledgeProcessingSpan, error)
|
||||
GetSpan(ctx context.Context, knowledgeID string, attempt int, spanID string) (*types.KnowledgeProcessingSpan, error)
|
||||
// CancelDescendants marks every descendant of a parent span as
|
||||
// "cancelled" with the given reason. Used by the tracker to
|
||||
// cascade an upstream failure across a stage's downstream subtree
|
||||
// without iterating in Go memory.
|
||||
CancelDescendants(ctx context.Context, knowledgeID string, attempt int, parentSpanID, reason string) (int64, error)
|
||||
// CancelAllOpenSpans flips every non-terminal (pending/running) span
|
||||
// for (knowledgeID, attempt) to "cancelled" in one statement,
|
||||
// regardless of tree position. Used by the user-cancel path where
|
||||
// fan-out stages (e.g. "多模态识别") flip themselves to done as soon
|
||||
// as they finish dispatching, while their async children are still
|
||||
// running — a tree walk that stops at terminal parents would miss
|
||||
// those orphan leaves.
|
||||
CancelAllOpenSpans(ctx context.Context, knowledgeID string, attempt int, errorCode, reason string) (int64, error)
|
||||
// CancelOpenSpansByName flips pending/running rows with the given span
|
||||
// name for (knowledgeID, attempt). Used before re-opening a subspan
|
||||
// after asynq retry or server restart so the trace tree does not
|
||||
// accumulate duplicate postprocess.summary / question rows.
|
||||
CancelOpenSpansByName(ctx context.Context, knowledgeID string, attempt int, name, errorCode, reason string) (int64, error)
|
||||
}
|
||||
|
||||
type knowledgeSpanRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewKnowledgeSpanRepository wires the GORM-backed implementation.
|
||||
func NewKnowledgeSpanRepository(db *gorm.DB) KnowledgeSpanRepository {
|
||||
return &knowledgeSpanRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *knowledgeSpanRepository) Upsert(ctx context.Context, row *types.KnowledgeProcessingSpan) error {
|
||||
if row == nil || row.KnowledgeID == "" || row.SpanID == "" {
|
||||
return errors.New("knowledgeSpanRepository.Upsert: knowledge_id and span_id required")
|
||||
}
|
||||
if row.Attempt == 0 {
|
||||
row.Attempt = 1
|
||||
}
|
||||
// Error fields may originate from third-party parser/model responses.
|
||||
// Normalize them at the persistence boundary so SQLite/PostgreSQL and the
|
||||
// tracing API never receive malformed UTF-8.
|
||||
row.ErrorCode = common.CleanInvalidUTF8(row.ErrorCode)
|
||||
row.ErrorMessage = common.CleanInvalidUTF8(row.ErrorMessage)
|
||||
row.ErrorDetail = common.CleanInvalidUTF8(row.ErrorDetail)
|
||||
// We let GORM populate created_at/updated_at via the autoCreate /
|
||||
// autoUpdate tags. ON CONFLICT updates only the fields that may
|
||||
// transition between calls — name/kind/parent are immutable once
|
||||
// set so we don't list them in DoUpdates (saves a few bytes per
|
||||
// write, and any mismatch indicates a programming error).
|
||||
//
|
||||
// CRITICAL: input / output / metadata are CONTENT fields that
|
||||
// individual call sites only fill when they have something to set.
|
||||
// EndSpan e.g. only sets `output`; if we always listed `input` in
|
||||
// DoUpdates, the End call would clobber the input set by Begin with
|
||||
// NULL. Same for metadata. Build the DoUpdates list dynamically and
|
||||
// skip these three columns when the incoming row has nothing to
|
||||
// write — so "no value" preserves the existing column instead of
|
||||
// nuking it.
|
||||
cols := []string{
|
||||
"status",
|
||||
"error_code",
|
||||
"error_message",
|
||||
"error_detail",
|
||||
"started_at",
|
||||
"finished_at",
|
||||
"duration_ms",
|
||||
"updated_at",
|
||||
}
|
||||
if row.Input != nil {
|
||||
cols = append(cols, "input")
|
||||
}
|
||||
if row.Output != nil {
|
||||
cols = append(cols, "output")
|
||||
}
|
||||
if row.Metadata != nil {
|
||||
cols = append(cols, "metadata")
|
||||
}
|
||||
return r.db.WithContext(ctx).Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{
|
||||
{Name: "knowledge_id"},
|
||||
{Name: "attempt"},
|
||||
{Name: "span_id"},
|
||||
},
|
||||
DoUpdates: clause.AssignmentColumns(cols),
|
||||
}).Create(row).Error
|
||||
}
|
||||
|
||||
func (r *knowledgeSpanRepository) NextAttempt(ctx context.Context, knowledgeID string) (int, error) {
|
||||
var max int
|
||||
err := r.db.WithContext(ctx).Model(&types.KnowledgeProcessingSpan{}).
|
||||
Where("knowledge_id = ?", knowledgeID).
|
||||
Select("COALESCE(MAX(attempt), 0)").
|
||||
Row().Scan(&max)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return max + 1, nil
|
||||
}
|
||||
|
||||
func (r *knowledgeSpanRepository) LatestAttempt(ctx context.Context, knowledgeID string) (int, error) {
|
||||
var max int
|
||||
err := r.db.WithContext(ctx).Model(&types.KnowledgeProcessingSpan{}).
|
||||
Where("knowledge_id = ?", knowledgeID).
|
||||
Select("COALESCE(MAX(attempt), 0)").
|
||||
Row().Scan(&max)
|
||||
return max, err
|
||||
}
|
||||
|
||||
func (r *knowledgeSpanRepository) ListByAttempt(ctx context.Context, knowledgeID string, attempt int) ([]types.KnowledgeProcessingSpan, error) {
|
||||
if knowledgeID == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var rows []types.KnowledgeProcessingSpan
|
||||
q := r.db.WithContext(ctx).Where("knowledge_id = ?", knowledgeID)
|
||||
if attempt > 0 {
|
||||
q = q.Where("attempt = ?", attempt)
|
||||
}
|
||||
// id ASC keeps the natural insertion order — useful for stable
|
||||
// rendering of fan-out subspans (e.g. multimodal.image[0..N] in
|
||||
// the order they were enqueued).
|
||||
err := q.Order("id ASC").Find(&rows).Error
|
||||
return rows, err
|
||||
}
|
||||
|
||||
func (r *knowledgeSpanRepository) GetSpan(ctx context.Context, knowledgeID string, attempt int, spanID string) (*types.KnowledgeProcessingSpan, error) {
|
||||
var row types.KnowledgeProcessingSpan
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("knowledge_id = ? AND attempt = ? AND span_id = ?", knowledgeID, attempt, spanID).
|
||||
Take(&row).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
// CancelDescendants performs an iterative SQL walk: each level we update
|
||||
// every row whose parent_span_id is in the previous level's span_id set,
|
||||
// flipping pending/running rows to cancelled. We bail when a level adds
|
||||
// zero rows (fixed point reached) or after a generous depth bound.
|
||||
//
|
||||
// Postgres-specific WITH RECURSIVE would be denser but harder to test on
|
||||
// the SQLite Lite backend. The iterative path stays portable.
|
||||
func (r *knowledgeSpanRepository) CancelDescendants(ctx context.Context, knowledgeID string, attempt int, parentSpanID, reason string) (int64, error) {
|
||||
reason = common.CleanInvalidUTF8(reason)
|
||||
frontier := []string{parentSpanID}
|
||||
var totalAffected int64
|
||||
for depth := 0; depth < 16 && len(frontier) > 0; depth++ {
|
||||
var nextFrontier []string
|
||||
// Find children of every span currently on the frontier
|
||||
// that are still in a non-terminal state — terminal rows
|
||||
// (done/failed/skipped/cancelled) are left as-is so the UI
|
||||
// can still see their original outcome.
|
||||
var children []types.KnowledgeProcessingSpan
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("knowledge_id = ? AND attempt = ? AND parent_span_id IN ? AND status IN ?",
|
||||
knowledgeID, attempt, frontier,
|
||||
[]string{types.SpanStatusPending, types.SpanStatusRunning}).
|
||||
Find(&children).Error
|
||||
if err != nil {
|
||||
return totalAffected, err
|
||||
}
|
||||
if len(children) == 0 {
|
||||
break
|
||||
}
|
||||
ids := make([]string, 0, len(children))
|
||||
for _, c := range children {
|
||||
ids = append(ids, c.SpanID)
|
||||
nextFrontier = append(nextFrontier, c.SpanID)
|
||||
}
|
||||
res := r.db.WithContext(ctx).Model(&types.KnowledgeProcessingSpan{}).
|
||||
Where("knowledge_id = ? AND attempt = ? AND span_id IN ?", knowledgeID, attempt, ids).
|
||||
Updates(map[string]any{
|
||||
"status": types.SpanStatusCancelled,
|
||||
"error_code": "UPSTREAM_FAILED",
|
||||
"error_message": reason,
|
||||
})
|
||||
if res.Error != nil {
|
||||
return totalAffected, res.Error
|
||||
}
|
||||
totalAffected += res.RowsAffected
|
||||
frontier = nextFrontier
|
||||
}
|
||||
return totalAffected, nil
|
||||
}
|
||||
|
||||
// CancelAllOpenSpans is the "abort the attempt" counterpart to
|
||||
// CancelDescendants. It avoids the BFS entirely so spans whose parent
|
||||
// is already terminal (typical for stage fan-outs that EndSpan as soon
|
||||
// as they finish dispatching async work) still get flipped to cancelled.
|
||||
// We deliberately do NOT touch finished_at / duration_ms here — the
|
||||
// span row remains observable in the trace tree with its original
|
||||
// start time and gets a cancelled status + reason, which is enough
|
||||
// for the UI to drop the running-bar styling.
|
||||
func (r *knowledgeSpanRepository) CancelAllOpenSpans(
|
||||
ctx context.Context, knowledgeID string, attempt int, errorCode, reason string,
|
||||
) (int64, error) {
|
||||
errorCode = common.CleanInvalidUTF8(errorCode)
|
||||
reason = common.CleanInvalidUTF8(reason)
|
||||
now := time.Now()
|
||||
updates := map[string]any{
|
||||
"status": types.SpanStatusCancelled,
|
||||
"error_code": errorCode,
|
||||
"error_message": reason,
|
||||
"finished_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
res := r.db.WithContext(ctx).Model(&types.KnowledgeProcessingSpan{}).
|
||||
Where("knowledge_id = ? AND attempt = ? AND status IN ?",
|
||||
knowledgeID, attempt,
|
||||
[]string{types.SpanStatusPending, types.SpanStatusRunning}).
|
||||
Updates(updates)
|
||||
if res.Error != nil {
|
||||
return 0, res.Error
|
||||
}
|
||||
return res.RowsAffected, nil
|
||||
}
|
||||
|
||||
func (r *knowledgeSpanRepository) CancelOpenSpansByName(
|
||||
ctx context.Context, knowledgeID string, attempt int, name, errorCode, reason string,
|
||||
) (int64, error) {
|
||||
if knowledgeID == "" || attempt <= 0 || name == "" {
|
||||
return 0, nil
|
||||
}
|
||||
errorCode = common.CleanInvalidUTF8(errorCode)
|
||||
reason = common.CleanInvalidUTF8(reason)
|
||||
now := time.Now()
|
||||
res := r.db.WithContext(ctx).Model(&types.KnowledgeProcessingSpan{}).
|
||||
Where("knowledge_id = ? AND attempt = ? AND name = ? AND status IN ?",
|
||||
knowledgeID, attempt, name,
|
||||
[]string{types.SpanStatusPending, types.SpanStatusRunning}).
|
||||
Updates(map[string]any{
|
||||
"status": types.SpanStatusCancelled,
|
||||
"error_code": errorCode,
|
||||
"error_message": reason,
|
||||
"finished_at": now,
|
||||
"updated_at": now,
|
||||
})
|
||||
if res.Error != nil {
|
||||
return 0, res.Error
|
||||
}
|
||||
return res.RowsAffected, nil
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// spansTestDDL mirrors migration 000053 for SQLite — same column order
|
||||
// minus the JSONB type (SQLite stores JSON as TEXT, the JSONMap Scanner
|
||||
// handles the round trip transparently). Inlined for the same reason
|
||||
// knowledgebase_sqlite_test.go inlines its DDL: GORM AutoMigrate doesn't
|
||||
// reproduce our PostgreSQL-flavoured schema cleanly.
|
||||
const spansTestDDL = `
|
||||
CREATE TABLE IF NOT EXISTS knowledge_processing_spans (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
knowledge_id VARCHAR(64) NOT NULL,
|
||||
attempt INTEGER NOT NULL DEFAULT 1,
|
||||
span_id VARCHAR(64) NOT NULL,
|
||||
parent_span_id VARCHAR(64),
|
||||
name VARCHAR(255) NOT NULL,
|
||||
kind VARCHAR(16) NOT NULL,
|
||||
status VARCHAR(16) NOT NULL,
|
||||
input TEXT,
|
||||
output TEXT,
|
||||
metadata TEXT,
|
||||
error_code VARCHAR(64),
|
||||
error_message TEXT,
|
||||
error_detail TEXT,
|
||||
started_at DATETIME,
|
||||
finished_at DATETIME,
|
||||
duration_ms BIGINT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
UNIQUE (knowledge_id, attempt, span_id)
|
||||
);
|
||||
`
|
||||
|
||||
func setupSpanTestRepo(t *testing.T) (KnowledgeSpanRepository, *gorm.DB) {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Exec(spansTestDDL).Error)
|
||||
return NewKnowledgeSpanRepository(db), db
|
||||
}
|
||||
|
||||
// TestKnowledgeSpanRepo_UpsertAndList covers the round-trip: a Begin
|
||||
// followed by an End for the same (kid, attempt, span_id) updates the
|
||||
// existing row in place, leaving exactly one row queryable by
|
||||
// ListByAttempt with the latest state.
|
||||
func TestKnowledgeSpanRepo_UpsertAndList(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
ctx := context.Background()
|
||||
kid := "kid-1"
|
||||
now := time.Now()
|
||||
row := &types.KnowledgeProcessingSpan{
|
||||
KnowledgeID: kid,
|
||||
Attempt: 1,
|
||||
SpanID: "span-A",
|
||||
Name: types.StageDocReader,
|
||||
Kind: types.SpanKindStage,
|
||||
Status: types.SpanStatusRunning,
|
||||
StartedAt: &now,
|
||||
}
|
||||
require.NoError(t, repo.Upsert(ctx, row))
|
||||
|
||||
// Second Upsert with same (kid, attempt, span_id) flips status and
|
||||
// sets finished_at — must overwrite, not insert a duplicate.
|
||||
finished := now.Add(2 * time.Second)
|
||||
row.Status = types.SpanStatusDone
|
||||
row.FinishedAt = &finished
|
||||
row.DurationMs = 2000
|
||||
require.NoError(t, repo.Upsert(ctx, row))
|
||||
|
||||
rows, err := repo.ListByAttempt(ctx, kid, 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, rows, 1, "Upsert must replace, not append")
|
||||
assert.Equal(t, types.SpanStatusDone, rows[0].Status)
|
||||
assert.Equal(t, int64(2000), rows[0].DurationMs)
|
||||
}
|
||||
|
||||
func TestKnowledgeSpanRepo_UpsertSanitizesErrorFields(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
now := time.Now()
|
||||
invalid := "prefix " + string([]byte{0xef, 0xbc, 0x2e}) + " suffix"
|
||||
|
||||
require.NoError(t, repo.Upsert(context.Background(), &types.KnowledgeProcessingSpan{
|
||||
KnowledgeID: "kid-invalid-utf8",
|
||||
Attempt: 1,
|
||||
SpanID: "span-invalid",
|
||||
Name: types.StageDocReader,
|
||||
Kind: types.SpanKindStage,
|
||||
Status: types.SpanStatusFailed,
|
||||
ErrorCode: invalid,
|
||||
ErrorMessage: invalid,
|
||||
ErrorDetail: invalid,
|
||||
StartedAt: &now,
|
||||
}))
|
||||
|
||||
rows, err := repo.ListByAttempt(context.Background(), "kid-invalid-utf8", 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, rows, 1)
|
||||
for _, value := range []string{rows[0].ErrorCode, rows[0].ErrorMessage, rows[0].ErrorDetail} {
|
||||
if !utf8.ValidString(value) {
|
||||
t.Fatalf("persisted error field is invalid UTF-8: % x", []byte(value))
|
||||
}
|
||||
if value != "prefix . suffix" {
|
||||
t.Fatalf("persisted error field = %q, want %q", value, "prefix . suffix")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestKnowledgeSpanRepo_NextAttempt confirms that NextAttempt allocates
|
||||
// a fresh number per knowledge, isolating reparse history. Critical
|
||||
// because the API layer renders attempt history by this number.
|
||||
func TestKnowledgeSpanRepo_NextAttempt(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
ctx := context.Background()
|
||||
kid := "kid-2"
|
||||
|
||||
a, err := repo.NextAttempt(ctx, kid)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, a, "first NextAttempt for fresh knowledge must be 1")
|
||||
|
||||
now := time.Now()
|
||||
require.NoError(t, repo.Upsert(ctx, &types.KnowledgeProcessingSpan{
|
||||
KnowledgeID: kid, Attempt: 1, SpanID: "root-1",
|
||||
Name: "knowledge_processing", Kind: types.SpanKindRoot,
|
||||
Status: types.SpanStatusRunning, StartedAt: &now,
|
||||
}))
|
||||
a, err = repo.NextAttempt(ctx, kid)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, a, "after one attempt exists, NextAttempt must be 2")
|
||||
|
||||
// Cross-knowledge isolation: a different kid stays at 1.
|
||||
other, err := repo.NextAttempt(ctx, "kid-other")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, other, "NextAttempt must scope to the knowledge_id")
|
||||
}
|
||||
|
||||
// TestKnowledgeSpanRepo_CancelDescendants verifies the cascade walk:
|
||||
// failing a stage cancels every pending/running descendant in its
|
||||
// subtree, while terminal states (done/skipped/failed) are left intact.
|
||||
func TestKnowledgeSpanRepo_CancelDescendants(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
ctx := context.Background()
|
||||
kid := "kid-3"
|
||||
now := time.Now()
|
||||
|
||||
// Tree: chunking → embedding (running) → batch[0] (running)
|
||||
// → multimodal (running) → image[0] (done)
|
||||
for _, r := range []*types.KnowledgeProcessingSpan{
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "chunking", Name: types.StageChunking, Kind: types.SpanKindStage, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "embedding", ParentSpanID: "chunking", Name: types.StageEmbedding, Kind: types.SpanKindStage, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "batch0", ParentSpanID: "embedding", Name: "embedding.batch[0]", Kind: types.SpanKindGeneration, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "multimodal", ParentSpanID: "chunking", Name: types.StageMultimodal, Kind: types.SpanKindStage, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "image0", ParentSpanID: "multimodal", Name: "multimodal.image[0]", Kind: types.SpanKindGeneration, Status: types.SpanStatusDone, StartedAt: &now},
|
||||
} {
|
||||
require.NoError(t, repo.Upsert(ctx, r))
|
||||
}
|
||||
|
||||
affected, err := repo.CancelDescendants(ctx, kid, 1, "chunking", "test reason")
|
||||
require.NoError(t, err)
|
||||
// Expected cancellations: embedding, batch0, multimodal (3 rows).
|
||||
// The done image0 is terminal and left alone.
|
||||
assert.Equal(t, int64(3), affected, "must cancel exactly the 3 pending/running descendants")
|
||||
|
||||
rows, err := repo.ListByAttempt(ctx, kid, 1)
|
||||
require.NoError(t, err)
|
||||
statusBy := map[string]string{}
|
||||
for _, r := range rows {
|
||||
statusBy[r.SpanID] = r.Status
|
||||
}
|
||||
assert.Equal(t, types.SpanStatusRunning, statusBy["chunking"], "the failed span itself stays untouched (FailSpan layer flips it)")
|
||||
assert.Equal(t, types.SpanStatusCancelled, statusBy["embedding"])
|
||||
assert.Equal(t, types.SpanStatusCancelled, statusBy["batch0"])
|
||||
assert.Equal(t, types.SpanStatusCancelled, statusBy["multimodal"])
|
||||
assert.Equal(t, types.SpanStatusDone, statusBy["image0"], "terminal states must not be touched")
|
||||
}
|
||||
|
||||
func TestKnowledgeSpanRepo_CancelOpenSpansByName(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
ctx := context.Background()
|
||||
kid := "kid-supersede"
|
||||
now := time.Now()
|
||||
|
||||
for _, r := range []*types.KnowledgeProcessingSpan{
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "sum-old", Name: "postprocess.summary", Kind: types.SpanKindSubSpan, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "sum-done", Name: "postprocess.summary", Kind: types.SpanKindSubSpan, Status: types.SpanStatusDone, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "q-old", Name: "postprocess.question", Kind: types.SpanKindSubSpan, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
} {
|
||||
require.NoError(t, repo.Upsert(ctx, r))
|
||||
}
|
||||
|
||||
affected, err := repo.CancelOpenSpansByName(ctx, kid, 1, "postprocess.summary", "TASK_SUPERSEDED", "retry")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), affected)
|
||||
|
||||
rows, err := repo.ListByAttempt(ctx, kid, 1)
|
||||
require.NoError(t, err)
|
||||
statusBy := map[string]string{}
|
||||
for _, r := range rows {
|
||||
statusBy[r.SpanID] = r.Status
|
||||
}
|
||||
assert.Equal(t, types.SpanStatusCancelled, statusBy["sum-old"])
|
||||
assert.Equal(t, types.SpanStatusDone, statusBy["sum-done"])
|
||||
assert.Equal(t, types.SpanStatusRunning, statusBy["q-old"])
|
||||
}
|
||||
|
||||
func TestKnowledgeSpanRepo_CancelPathsSanitizeErrorFields(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
ctx := context.Background()
|
||||
kid := "kid-cancel-invalid-utf8"
|
||||
now := time.Now()
|
||||
for _, r := range []*types.KnowledgeProcessingSpan{
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "root", Name: "root", Kind: types.SpanKindRoot, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "descendant", ParentSpanID: "root", Name: "descendant", Kind: types.SpanKindSubSpan, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "bulk", Name: "bulk", Kind: types.SpanKindSubSpan, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
{KnowledgeID: kid, Attempt: 1, SpanID: "named", Name: "named", Kind: types.SpanKindSubSpan, Status: types.SpanStatusRunning, StartedAt: &now},
|
||||
} {
|
||||
require.NoError(t, repo.Upsert(ctx, r))
|
||||
}
|
||||
|
||||
invalidCode := "CODE " + string([]byte{0xef, 0xbc, 0x2e})
|
||||
invalidReason := "reason " + string([]byte{0xef, 0xbc, 0x2e})
|
||||
_, err := repo.CancelOpenSpansByName(ctx, kid, 1, "named", invalidCode, invalidReason)
|
||||
require.NoError(t, err)
|
||||
_, err = repo.CancelDescendants(ctx, kid, 1, "root", invalidReason)
|
||||
require.NoError(t, err)
|
||||
_, err = repo.CancelAllOpenSpans(ctx, kid, 1, invalidCode, invalidReason)
|
||||
require.NoError(t, err)
|
||||
|
||||
rows, err := repo.ListByAttempt(ctx, kid, 1)
|
||||
require.NoError(t, err)
|
||||
byID := make(map[string]types.KnowledgeProcessingSpan, len(rows))
|
||||
for _, row := range rows {
|
||||
byID[row.SpanID] = row
|
||||
}
|
||||
assert.Equal(t, "UPSTREAM_FAILED", byID["descendant"].ErrorCode)
|
||||
assert.Equal(t, "reason .", byID["descendant"].ErrorMessage)
|
||||
for _, id := range []string{"root", "bulk", "named"} {
|
||||
assert.Equal(t, "CODE .", byID[id].ErrorCode, id)
|
||||
assert.Equal(t, "reason .", byID[id].ErrorMessage, id)
|
||||
assert.True(t, utf8.ValidString(byID[id].ErrorCode), id)
|
||||
assert.True(t, utf8.ValidString(byID[id].ErrorMessage), id)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKnowledgeSpanRepo_ListAttemptIsolation guarantees that different
|
||||
// attempts of the same knowledge stay queryable independently — the
|
||||
// foundation for the "show history" UI navigation (?attempt=N).
|
||||
func TestKnowledgeSpanRepo_ListAttemptIsolation(t *testing.T) {
|
||||
repo, _ := setupSpanTestRepo(t)
|
||||
ctx := context.Background()
|
||||
kid := "kid-history"
|
||||
now := time.Now()
|
||||
|
||||
for _, attempt := range []int{1, 2} {
|
||||
require.NoError(t, repo.Upsert(ctx, &types.KnowledgeProcessingSpan{
|
||||
KnowledgeID: kid, Attempt: attempt, SpanID: "root",
|
||||
Name: "knowledge_processing", Kind: types.SpanKindRoot,
|
||||
Status: types.SpanStatusDone, StartedAt: &now,
|
||||
}))
|
||||
}
|
||||
a1, err := repo.ListByAttempt(ctx, kid, 1)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, a1, 1)
|
||||
a2, err := repo.ListByAttempt(ctx, kid, 2)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, a2, 1)
|
||||
|
||||
all, err := repo.ListByAttempt(ctx, kid, 0)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, all, 2, "attempt=0 returns all attempts (used by housekeeping)")
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// SetKnowledgeTags replaces all tags for a single knowledge entry.
|
||||
// It deletes existing relations and inserts new ones in a transaction.
|
||||
func (r *knowledgeRepository) SetKnowledgeTags(
|
||||
ctx context.Context,
|
||||
knowledgeID string,
|
||||
tagIDs []string,
|
||||
) error {
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var knowledge types.Knowledge
|
||||
if err := tx.Select("id", "tenant_id", "knowledge_base_id").First(&knowledge, "id = ?", knowledgeID).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := lockActiveKnowledgeBase(tx, knowledge.TenantID, knowledge.KnowledgeBaseID); err != nil {
|
||||
return err
|
||||
}
|
||||
// Delete all existing tag relations for this knowledge
|
||||
if err := tx.Where("knowledge_id = ?", knowledgeID).
|
||||
Delete(&types.KnowledgeTagRelation{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// Insert new relations (skip empty and duplicate IDs)
|
||||
if len(tagIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{}, len(tagIDs))
|
||||
now := time.Now()
|
||||
relations := make([]types.KnowledgeTagRelation, 0, len(tagIDs))
|
||||
for _, tagID := range tagIDs {
|
||||
if tagID == "" {
|
||||
continue
|
||||
}
|
||||
if _, dup := seen[tagID]; dup {
|
||||
continue
|
||||
}
|
||||
seen[tagID] = struct{}{}
|
||||
relations = append(relations, types.KnowledgeTagRelation{
|
||||
KnowledgeID: knowledgeID,
|
||||
TagID: tagID,
|
||||
CreatedAt: now,
|
||||
})
|
||||
}
|
||||
if len(relations) == 0 {
|
||||
return nil
|
||||
}
|
||||
return tx.Create(&relations).Error
|
||||
})
|
||||
}
|
||||
|
||||
// AddKnowledgeTagRelations incrementally adds validated relations and never
|
||||
// removes existing manual tags. The composite primary key plus DO NOTHING
|
||||
// makes retries and duplicate deliveries idempotent.
|
||||
func (r *knowledgeRepository) AddKnowledgeTagRelations(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
kbID, knowledgeID string,
|
||||
tagIDs []string,
|
||||
) error {
|
||||
seen := make(map[string]struct{}, len(tagIDs))
|
||||
unique := make([]string, 0, len(tagIDs))
|
||||
for _, id := range tagIDs {
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[id]; ok {
|
||||
continue
|
||||
}
|
||||
seen[id] = struct{}{}
|
||||
unique = append(unique, id)
|
||||
}
|
||||
if len(unique) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := lockActiveKnowledgeBase(tx, tenantID, kbID); err != nil {
|
||||
return err
|
||||
}
|
||||
var knowledgeCount int64
|
||||
if err := tx.Model(&types.Knowledge{}).
|
||||
Where("id = ? AND tenant_id = ? AND knowledge_base_id = ? AND deleted_at IS NULL", knowledgeID, tenantID, kbID).
|
||||
Where("parse_status NOT IN ?", []string{
|
||||
types.ParseStatusCancelled,
|
||||
types.ParseStatusDeleting,
|
||||
types.ParseStatusFailed,
|
||||
}).
|
||||
Count(&knowledgeCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if knowledgeCount != 1 {
|
||||
return fmt.Errorf("knowledge %s does not belong to tenant %d and knowledge base %s", knowledgeID, tenantID, kbID)
|
||||
}
|
||||
|
||||
var tagCount int64
|
||||
if err := tx.Model(&types.KnowledgeTag{}).
|
||||
Where("tenant_id = ? AND knowledge_base_id = ? AND id IN ?", tenantID, kbID, unique).
|
||||
Count(&tagCount).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if tagCount != int64(len(unique)) {
|
||||
return fmt.Errorf("one or more tags do not belong to tenant %d and knowledge base %s", tenantID, kbID)
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
relations := make([]types.KnowledgeTagRelation, 0, len(unique))
|
||||
for _, tagID := range unique {
|
||||
relations = append(relations, types.KnowledgeTagRelation{
|
||||
KnowledgeID: knowledgeID,
|
||||
TagID: tagID,
|
||||
CreatedAt: now,
|
||||
})
|
||||
}
|
||||
return tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&relations).Error
|
||||
})
|
||||
}
|
||||
|
||||
// GetKnowledgeTags returns tags for multiple knowledge IDs.
|
||||
// The result is a map from knowledge ID to its tag list.
|
||||
func (r *knowledgeRepository) GetKnowledgeTags(
|
||||
ctx context.Context,
|
||||
knowledgeIDs []string,
|
||||
) (map[string][]*types.KnowledgeTag, error) {
|
||||
result := make(map[string][]*types.KnowledgeTag)
|
||||
if len(knowledgeIDs) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// Query relations and join with knowledge_tags to get full tag info
|
||||
type relationWithTag struct {
|
||||
KnowledgeID string `gorm:"column:knowledge_id"`
|
||||
types.KnowledgeTag
|
||||
}
|
||||
var rows []relationWithTag
|
||||
if err := r.db.WithContext(ctx).
|
||||
Table("knowledge_tag_relations AS ktr").
|
||||
Select("ktr.knowledge_id, kt.id, kt.seq_id, kt.tenant_id, kt.knowledge_base_id, kt.name, kt.color, kt.sort_order, kt.created_at, kt.updated_at").
|
||||
Joins("JOIN knowledge_tags AS kt ON ktr.tag_id = kt.id").
|
||||
Where("ktr.knowledge_id IN (?)", knowledgeIDs).
|
||||
Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
for _, row := range rows {
|
||||
tag := row.KnowledgeTag
|
||||
result[row.KnowledgeID] = append(result[row.KnowledgeID], &tag)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// DeleteKnowledgeTagRelations deletes all tag relations for a knowledge entry.
|
||||
func (r *knowledgeRepository) DeleteKnowledgeTagRelations(
|
||||
ctx context.Context,
|
||||
knowledgeID string,
|
||||
) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("knowledge_id = ?", knowledgeID).
|
||||
Delete(&types.KnowledgeTagRelation{}).Error
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const knowledgeTagTestDDL = `
|
||||
CREATE TABLE IF NOT EXISTS knowledges (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
knowledge_base_id VARCHAR(36) NOT NULL,
|
||||
type VARCHAR(50) NOT NULL DEFAULT 'file',
|
||||
title VARCHAR(255) NOT NULL DEFAULT '',
|
||||
parse_status VARCHAR(50) NOT NULL DEFAULT 'completed',
|
||||
deleted_at DATETIME
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS knowledge_tags (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
seq_id INTEGER NOT NULL,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
knowledge_base_id VARCHAR(36) NOT NULL,
|
||||
name VARCHAR(128) NOT NULL,
|
||||
color VARCHAR(32),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS knowledge_tag_relations (
|
||||
knowledge_id VARCHAR(36) NOT NULL,
|
||||
tag_id VARCHAR(36) NOT NULL,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (knowledge_id, tag_id)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS chunks (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
knowledge_base_id VARCHAR(36) NOT NULL,
|
||||
tag_id VARCHAR(36),
|
||||
deleted_at DATETIME
|
||||
);
|
||||
`
|
||||
|
||||
func setupKnowledgeTagTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db := setupKnowledgeTestDB(t)
|
||||
require.NoError(t, db.Exec(knowledgeBasesTestDDL).Error)
|
||||
require.NoError(t, db.Exec(knowledgeTagTestDDL).Error)
|
||||
return db
|
||||
}
|
||||
|
||||
func seedKnowledgeTagFixture(t *testing.T, db *gorm.DB) (kbID, knowledgeID, tagA, tagB string) {
|
||||
t.Helper()
|
||||
kbID = uuid.New().String()
|
||||
kb := makeKB(nil)
|
||||
kb.ID = kbID
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
knowledgeID = uuid.New().String()
|
||||
tagA = uuid.New().String()
|
||||
tagB = uuid.New().String()
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, parse_status)
|
||||
VALUES (?, 1, ?, 'file', 'doc', 'completed')
|
||||
`, knowledgeID, kbID).Error)
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledge_tags (id, seq_id, tenant_id, knowledge_base_id, name)
|
||||
VALUES (?, 1, 1, ?, 'alpha'), (?, 2, 1, ?, 'beta')
|
||||
`, tagA, kbID, tagB, kbID).Error)
|
||||
return kbID, knowledgeID, tagA, tagB
|
||||
}
|
||||
|
||||
func TestSetKnowledgeTags_DedupesAndReplaces(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
repo := &knowledgeRepository{db: db}
|
||||
_, knowledgeID, tagA, tagB := seedKnowledgeTagFixture(t, db)
|
||||
ctx := context.Background()
|
||||
|
||||
err := repo.SetKnowledgeTags(ctx, knowledgeID, []string{tagA, tagA, tagB})
|
||||
require.NoError(t, err)
|
||||
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.KnowledgeTagRelation{}).
|
||||
Where("knowledge_id = ?", knowledgeID).Count(&count).Error)
|
||||
assert.Equal(t, int64(2), count)
|
||||
|
||||
err = repo.SetKnowledgeTags(ctx, knowledgeID, []string{tagB})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Model(&types.KnowledgeTagRelation{}).
|
||||
Where("knowledge_id = ?", knowledgeID).Count(&count).Error)
|
||||
assert.Equal(t, int64(1), count)
|
||||
|
||||
err = repo.SetKnowledgeTags(ctx, knowledgeID, nil)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Model(&types.KnowledgeTagRelation{}).
|
||||
Where("knowledge_id = ?", knowledgeID).Count(&count).Error)
|
||||
assert.Equal(t, int64(0), count)
|
||||
}
|
||||
|
||||
func TestAddKnowledgeTagRelations_IsIncrementalIdempotentAndScoped(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
repo := &knowledgeRepository{db: db}
|
||||
kbID, knowledgeID, tagA, tagB := seedKnowledgeTagFixture(t, db)
|
||||
ctx := context.Background()
|
||||
|
||||
require.NoError(t, repo.SetKnowledgeTags(ctx, knowledgeID, []string{tagA}))
|
||||
require.NoError(t, repo.AddKnowledgeTagRelations(ctx, 1, kbID, knowledgeID, []string{tagA, tagB, tagB}))
|
||||
require.NoError(t, repo.AddKnowledgeTagRelations(ctx, 1, kbID, knowledgeID, []string{tagB}))
|
||||
|
||||
var relations []types.KnowledgeTagRelation
|
||||
require.NoError(t, db.Where("knowledge_id = ?", knowledgeID).Find(&relations).Error)
|
||||
assert.Len(t, relations, 2)
|
||||
|
||||
otherKBTag := uuid.New().String()
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledge_tags (id, seq_id, tenant_id, knowledge_base_id, name)
|
||||
VALUES (?, 3, 1, ?, 'other')
|
||||
`, otherKBTag, uuid.New().String()).Error)
|
||||
err := repo.AddKnowledgeTagRelations(ctx, 1, kbID, knowledgeID, []string{otherKBTag})
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestAddKnowledgeTagRelations_RejectsCancelledKnowledge(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
repo := &knowledgeRepository{db: db}
|
||||
kbID, knowledgeID, _, tagB := seedKnowledgeTagFixture(t, db)
|
||||
require.NoError(t, db.Model(&types.Knowledge{}).
|
||||
Where("id = ?", knowledgeID).
|
||||
Update("parse_status", types.ParseStatusCancelled).Error)
|
||||
|
||||
err := repo.AddKnowledgeTagRelations(context.Background(), 1, kbID, knowledgeID, []string{tagB})
|
||||
require.Error(t, err)
|
||||
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.KnowledgeTagRelation{}).
|
||||
Where("knowledge_id = ?", knowledgeID).
|
||||
Count(&count).Error)
|
||||
assert.Zero(t, count)
|
||||
}
|
||||
|
||||
func TestGetKnowledgeTags_ReturnsTagDetails(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
repo := &knowledgeRepository{db: db}
|
||||
_, knowledgeID, tagA, tagB := seedKnowledgeTagFixture(t, db)
|
||||
ctx := context.Background()
|
||||
|
||||
require.NoError(t, repo.SetKnowledgeTags(ctx, knowledgeID, []string{tagA, tagB}))
|
||||
|
||||
tagMap, err := repo.GetKnowledgeTags(ctx, []string{knowledgeID})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, tagMap[knowledgeID], 2)
|
||||
|
||||
names := map[string]bool{}
|
||||
for _, tag := range tagMap[knowledgeID] {
|
||||
names[tag.Name] = true
|
||||
}
|
||||
assert.True(t, names["alpha"])
|
||||
assert.True(t, names["beta"])
|
||||
}
|
||||
|
||||
func TestDeleteKnowledgeTagRelations(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
repo := &knowledgeRepository{db: db}
|
||||
_, knowledgeID, tagA, _ := seedKnowledgeTagFixture(t, db)
|
||||
ctx := context.Background()
|
||||
|
||||
require.NoError(t, repo.SetKnowledgeTags(ctx, knowledgeID, []string{tagA}))
|
||||
require.NoError(t, repo.DeleteKnowledgeTagRelations(ctx, knowledgeID))
|
||||
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.KnowledgeTagRelation{}).
|
||||
Where("knowledge_id = ?", knowledgeID).Count(&count).Error)
|
||||
assert.Equal(t, int64(0), count)
|
||||
}
|
||||
|
||||
func TestApplyKnowledgeListFilter_TagIDsOrSemantics(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
repo := &knowledgeRepository{db: db}
|
||||
ctx := context.Background()
|
||||
|
||||
kbID := uuid.New().String()
|
||||
kb := makeKB(nil)
|
||||
kb.ID = kbID
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
docA := uuid.New().String()
|
||||
docB := uuid.New().String()
|
||||
docC := uuid.New().String()
|
||||
tagA := uuid.New().String()
|
||||
tagB := uuid.New().String()
|
||||
|
||||
for _, row := range []struct{ id, title string }{
|
||||
{docA, "a"}, {docB, "b"}, {docC, "c"},
|
||||
} {
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, parse_status)
|
||||
VALUES (?, 1, ?, 'file', ?, 'completed')
|
||||
`, row.id, kbID, row.title).Error)
|
||||
}
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledge_tags (id, seq_id, tenant_id, knowledge_base_id, name)
|
||||
VALUES (?, 1, 1, ?, 't-a'), (?, 2, 1, ?, 't-b')
|
||||
`, tagA, kbID, tagB, kbID).Error)
|
||||
require.NoError(t, repo.SetKnowledgeTags(ctx, docA, []string{tagA}))
|
||||
require.NoError(t, repo.SetKnowledgeTags(ctx, docB, []string{tagB}))
|
||||
require.NoError(t, repo.SetKnowledgeTags(ctx, docC, []string{tagA, tagB}))
|
||||
|
||||
query := db.WithContext(ctx).Model(&types.Knowledge{}).
|
||||
Where("tenant_id = ? AND knowledge_base_id = ?", uint64(1), kbID)
|
||||
query = applyKnowledgeListFilter(query, types.KnowledgeListFilter{
|
||||
TagIDs: []string{tagA, tagB},
|
||||
})
|
||||
|
||||
var ids []string
|
||||
require.NoError(t, query.Pluck("id", &ids).Error)
|
||||
assert.ElementsMatch(t, []string{docA, docB, docC}, ids)
|
||||
}
|
||||
|
||||
func TestBatchCountReferences_ScopedToKnowledgeBase(t *testing.T) {
|
||||
db := setupKnowledgeTagTestDB(t)
|
||||
knowledgeRepo := &knowledgeRepository{db: db}
|
||||
tagRepo := &knowledgeTagRepository{db: db}
|
||||
ctx := context.Background()
|
||||
|
||||
kb1 := uuid.New().String()
|
||||
kb2 := uuid.New().String()
|
||||
for _, id := range []string{kb1, kb2} {
|
||||
kb := makeKB(nil)
|
||||
kb.ID = id
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
}
|
||||
doc1 := uuid.New().String()
|
||||
doc2 := uuid.New().String()
|
||||
tag1 := uuid.New().String()
|
||||
tag2 := uuid.New().String()
|
||||
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledges (id, tenant_id, knowledge_base_id, type, title, parse_status)
|
||||
VALUES (?, 1, ?, 'file', 'kb1-doc', 'completed'), (?, 1, ?, 'file', 'kb2-doc', 'completed')
|
||||
`, doc1, kb1, doc2, kb2).Error)
|
||||
require.NoError(t, db.Exec(`
|
||||
INSERT INTO knowledge_tags (id, seq_id, tenant_id, knowledge_base_id, name)
|
||||
VALUES (?, 1, 1, ?, 'tag-kb1'), (?, 2, 1, ?, 'tag-kb2')
|
||||
`, tag1, kb1, tag2, kb2).Error)
|
||||
require.NoError(t, knowledgeRepo.SetKnowledgeTags(ctx, doc1, []string{tag1}))
|
||||
// Stale relation: doc in kb2 still linked to tag1 (simulates pre-fix move bug)
|
||||
require.NoError(t, knowledgeRepo.SetKnowledgeTags(ctx, doc2, []string{tag1}))
|
||||
|
||||
counts, err := tagRepo.BatchCountReferences(ctx, 1, kb1, []string{tag1})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), counts[tag1].KnowledgeCount)
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/common"
|
||||
"github.com/Tencent/WeKnora/internal/knowledgeversion"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// CreateVersionKnowledge keeps staging documents and their accounted bytes
|
||||
// atomic so an interrupted copy can later reconcile storage exactly once.
|
||||
func (r *knowledgeRepository) CreateVersionKnowledge(ctx context.Context, doc *types.Knowledge) error {
|
||||
tenant, ok := types.TenantIDFromContext(ctx)
|
||||
if !knowledgeversion.IsWorker(ctx) || !ok || doc.TenantID != tenant || doc.StorageSize < 0 {
|
||||
return fmt.Errorf("version document requires its owning workspace worker")
|
||||
}
|
||||
doc.ErrorMessage = common.CleanInvalidUTF8(doc.ErrorMessage)
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Create(doc).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
result := tx.Model(&types.Tenant{}).Where("id = ?", tenant).UpdateColumn("storage_used", gorm.Expr("storage_used + ?", doc.StorageSize))
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return fmt.Errorf("version workspace not found")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/knowledgeversion"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// CopyVersionGraph copies imported PostgreSQL evidence without rerunning AI
|
||||
// extraction. Provenance remains immutable; only proven native pointers change.
|
||||
func (r *knowledgeRepository) CopyVersionGraph(ctx context.Context, source, target string, documents, chunks map[string]string) error {
|
||||
tenant, ok := types.TenantIDFromContext(ctx)
|
||||
if !knowledgeversion.IsWorker(ctx) || !ok || tenant == 0 || source == target {
|
||||
return fmt.Errorf("graph version copy requires an isolated workspace worker")
|
||||
}
|
||||
db := r.db.WithContext(ctx)
|
||||
if !knowledgeversion.GraphTablesExist(db) {
|
||||
return nil
|
||||
}
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
var bases []types.KnowledgeBase
|
||||
if err := tx.Where("tenant_id=? AND id IN ?", tenant, []string{source, target}).Find(&bases).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
validTarget := false
|
||||
for _, kb := range bases {
|
||||
if kb.ID == target && kb.IsTemporary && kb.VersionState == "copying" {
|
||||
validTarget = true
|
||||
}
|
||||
}
|
||||
if len(bases) != 2 || !validTarget {
|
||||
return fmt.Errorf("graph version target must be a temporary workspace copy")
|
||||
}
|
||||
for _, model := range []any{&knowledgeversion.GraphRecord{}, &knowledgeversion.GraphSource{}} {
|
||||
var count int64
|
||||
if err := tx.Model(model).Where("tenant_id=? AND knowledge_base_id=?", tenant, target).Count(&count).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if count != 0 {
|
||||
return fmt.Errorf("graph version target must be empty")
|
||||
}
|
||||
}
|
||||
for offset := 0; ; offset += 200 {
|
||||
var records []knowledgeversion.GraphRecord
|
||||
if err := tx.Where("tenant_id=? AND knowledge_base_id=?", tenant, source).Order("source_table,legacy_id").Limit(200).Offset(offset).Find(&records).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(records) == 0 {
|
||||
break
|
||||
}
|
||||
for i := range records {
|
||||
records[i].KnowledgeBaseID = target
|
||||
}
|
||||
if err := tx.Create(&records).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for offset := 0; ; offset += 200 {
|
||||
var evidence []knowledgeversion.GraphSource
|
||||
if err := tx.Where("tenant_id=? AND knowledge_base_id=?", tenant, source).Order("source_chunk_id").Limit(200).Offset(offset).Find(&evidence).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(evidence) == 0 {
|
||||
break
|
||||
}
|
||||
ids := []string{}
|
||||
for _, e := range evidence {
|
||||
if e.NativeChunkID != nil {
|
||||
ids = append(ids, *e.NativeChunkID)
|
||||
if dst := chunks[*e.NativeChunkID]; dst != "" {
|
||||
ids = append(ids, dst)
|
||||
}
|
||||
}
|
||||
}
|
||||
loaded := []types.Chunk{}
|
||||
if len(ids) > 0 {
|
||||
if err := tx.Where("tenant_id=? AND knowledge_base_id IN ? AND id IN ?", tenant, []string{source, target}, ids).Find(&loaded).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
byID := map[string]types.Chunk{}
|
||||
for _, c := range loaded {
|
||||
byID[c.ID] = c
|
||||
}
|
||||
for i := range evidence {
|
||||
e := &evidence[i]
|
||||
e.KnowledgeBaseID = target
|
||||
oldID, oldDoc, oldHash := e.NativeChunkID, e.NativeKnowledgeID, e.NativeContentSHA256
|
||||
e.NativeChunkID, e.NativeKnowledgeID, e.NativeContentSHA256 = nil, nil, nil
|
||||
if oldID == nil {
|
||||
continue
|
||||
}
|
||||
e.MappingState = "stale"
|
||||
src, srcOK := byID[*oldID]
|
||||
dst, dstOK := byID[chunks[*oldID]]
|
||||
if !srcOK || !dstOK || oldDoc == nil || oldHash == nil || src.KnowledgeBaseID != source || dst.KnowledgeBaseID != target || src.KnowledgeID != *oldDoc || documents[*oldDoc] != dst.KnowledgeID {
|
||||
continue
|
||||
}
|
||||
sourceHash := sha256.Sum256([]byte(src.Content))
|
||||
if hex.EncodeToString(sourceHash[:]) != *oldHash {
|
||||
continue
|
||||
}
|
||||
targetHash := sha256.Sum256([]byte(dst.Content))
|
||||
hash := hex.EncodeToString(targetHash[:])
|
||||
e.NativeChunkID, e.NativeKnowledgeID, e.NativeContentSHA256 = &dst.ID, &dst.KnowledgeID, &hash
|
||||
e.MappingState = "exact_content_hash"
|
||||
}
|
||||
if err := tx.Create(&evidence).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,310 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
var ErrKnowledgeBaseNotFound = errors.New("knowledge base not found")
|
||||
|
||||
// knowledgeBaseRepository implements the KnowledgeBaseRepository interface
|
||||
type knowledgeBaseRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewKnowledgeBaseRepository creates a new knowledge base repository
|
||||
func NewKnowledgeBaseRepository(db *gorm.DB) interfaces.KnowledgeBaseRepository {
|
||||
return &knowledgeBaseRepository{db: db}
|
||||
}
|
||||
|
||||
// CreateKnowledgeBase creates a new knowledge base
|
||||
func (r *knowledgeBaseRepository) CreateKnowledgeBase(ctx context.Context, kb *types.KnowledgeBase) error {
|
||||
return r.db.WithContext(ctx).Create(kb).Error
|
||||
}
|
||||
|
||||
// GetKnowledgeBaseByID gets a knowledge base by id (no tenant scope; caller must enforce isolation where needed)
|
||||
func (r *knowledgeBaseRepository) GetKnowledgeBaseByID(ctx context.Context, id string) (*types.KnowledgeBase, error) {
|
||||
var kb types.KnowledgeBase
|
||||
if err := r.db.WithContext(ctx).Where("id = ?", id).First(&kb).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrKnowledgeBaseNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &kb, nil
|
||||
}
|
||||
|
||||
// GetKnowledgeBaseByIDAndTenant gets a knowledge base by id only if it belongs to the given tenant (enforces tenant isolation)
|
||||
func (r *knowledgeBaseRepository) GetKnowledgeBaseByIDAndTenant(ctx context.Context, id string, tenantID uint64) (*types.KnowledgeBase, error) {
|
||||
var kb types.KnowledgeBase
|
||||
if err := r.db.WithContext(ctx).Where("id = ? AND tenant_id = ?", id, tenantID).First(&kb).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrKnowledgeBaseNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &kb, nil
|
||||
}
|
||||
|
||||
// GetKnowledgeBaseByIDs gets knowledge bases by multiple ids
|
||||
func (r *knowledgeBaseRepository) GetKnowledgeBaseByIDs(ctx context.Context, ids []string) ([]*types.KnowledgeBase, error) {
|
||||
if len(ids) == 0 {
|
||||
return []*types.KnowledgeBase{}, nil
|
||||
}
|
||||
var kbs []*types.KnowledgeBase
|
||||
if err := r.db.WithContext(ctx).Where("id IN ?", ids).Find(&kbs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return kbs, nil
|
||||
}
|
||||
|
||||
// ListKnowledgeBases lists all knowledge bases
|
||||
func (r *knowledgeBaseRepository) ListKnowledgeBases(ctx context.Context) ([]*types.KnowledgeBase, error) {
|
||||
var kbs []*types.KnowledgeBase
|
||||
if err := r.db.WithContext(ctx).Find(&kbs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return kbs, nil
|
||||
}
|
||||
|
||||
// ListKnowledgeBasesByTenantID lists all knowledge bases by tenant id.
|
||||
//
|
||||
// Ordering used to also include `is_pinned DESC, pinned_at DESC` so the
|
||||
// repository would return tenant-wide pinned rows first. That column is
|
||||
// no longer the source of truth (see migration 000050) — pin state is
|
||||
// now per (user, kb) and applied by the service layer after enrichment.
|
||||
// We keep `created_at DESC` here so callers that don't enrich (chat
|
||||
// pipeline, agent editor, IM commands) still get a stable ordering.
|
||||
func (r *knowledgeBaseRepository) ListKnowledgeBasesByTenantID(
|
||||
ctx context.Context, tenantID uint64,
|
||||
) ([]*types.KnowledgeBase, error) {
|
||||
var kbs []*types.KnowledgeBase
|
||||
if err := r.db.WithContext(ctx).Where("tenant_id = ? AND is_temporary = ?", tenantID, false).
|
||||
Order("created_at DESC").Find(&kbs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return kbs, nil
|
||||
}
|
||||
|
||||
// userKBPinRow mirrors the user_kb_pins table. Kept local to the
|
||||
// repository because it never escapes the package; callers see the
|
||||
// higher-level map[kb_id]pinned_at returned by ListUserKBPinIDs.
|
||||
type userKBPinRow struct {
|
||||
TenantID uint64 `gorm:"column:tenant_id"`
|
||||
UserID string `gorm:"column:user_id"`
|
||||
KBID string `gorm:"column:kb_id"`
|
||||
PinnedAt time.Time `gorm:"column:pinned_at"`
|
||||
}
|
||||
|
||||
func (userKBPinRow) TableName() string { return "user_kb_pins" }
|
||||
|
||||
// SetUserKBPin upserts (pinned=true) or deletes (pinned=false) the row
|
||||
// for the given (tenant, user, kb) triple. The returned pinned_at is
|
||||
// nil when pinned=false; otherwise it carries the timestamp written
|
||||
// to the row (either the existing one if the row already existed, or
|
||||
// the current time on insert) so the caller can stamp the response
|
||||
// without a follow-up SELECT.
|
||||
func (r *knowledgeBaseRepository) SetUserKBPin(
|
||||
ctx context.Context, tenantID uint64, userID string, kbID string, pinned bool,
|
||||
) (*time.Time, error) {
|
||||
if userID == "" {
|
||||
return nil, errors.New("user_kb_pins: empty user_id")
|
||||
}
|
||||
if !pinned {
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND user_id = ? AND kb_id = ?", tenantID, userID, kbID).
|
||||
Delete(&userKBPinRow{}).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Upsert with idempotent INSERT … ON CONFLICT DO NOTHING. We then
|
||||
// SELECT to learn whether an existing row's pinned_at survived (so
|
||||
// repeated calls return a stable timestamp instead of bumping it).
|
||||
row := userKBPinRow{
|
||||
TenantID: tenantID,
|
||||
UserID: userID,
|
||||
KBID: kbID,
|
||||
PinnedAt: time.Now(),
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND user_id = ? AND kb_id = ?", tenantID, userID, kbID).
|
||||
Attrs(userKBPinRow{PinnedAt: row.PinnedAt}).
|
||||
FirstOrCreate(&row).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pa := row.PinnedAt
|
||||
return &pa, nil
|
||||
}
|
||||
|
||||
// ListUserKBPinIDs returns every KB id this user has personally pinned
|
||||
// in this tenant, mapped to its pinned_at. Returns an empty map (not
|
||||
// nil) when there are no pins, so callers can do `len(m) == 0` checks
|
||||
// without a nil guard.
|
||||
func (r *knowledgeBaseRepository) ListUserKBPinIDs(
|
||||
ctx context.Context, tenantID uint64, userID string,
|
||||
) (map[string]time.Time, error) {
|
||||
out := make(map[string]time.Time)
|
||||
if userID == "" {
|
||||
return out, nil
|
||||
}
|
||||
var rows []userKBPinRow
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND user_id = ?", tenantID, userID).
|
||||
Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
out[row.KBID] = row.PinnedAt
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// UpdateKnowledgeBase updates a knowledge base
|
||||
func (r *knowledgeBaseRepository) UpdateKnowledgeBase(ctx context.Context, kb *types.KnowledgeBase) error {
|
||||
return r.db.WithContext(ctx).Save(kb).Error
|
||||
}
|
||||
|
||||
// DeleteKnowledgeBase atomically hides the KB and removes its categories.
|
||||
func (r *knowledgeBaseRepository) DeleteKnowledgeBase(ctx context.Context, id string) error {
|
||||
return r.deleteKnowledgeBaseCategories(ctx, id, false)
|
||||
}
|
||||
|
||||
// CleanupDeletedKnowledgeBaseCategories also handles retries with no live documents.
|
||||
func (r *knowledgeBaseRepository) CleanupDeletedKnowledgeBaseCategories(ctx context.Context, id string) error {
|
||||
return r.deleteKnowledgeBaseCategories(ctx, id, true)
|
||||
}
|
||||
|
||||
// Category writers share the delete lock so in-flight requests cannot resurrect
|
||||
// categories or relations after a whole-KB deletion commits.
|
||||
func lockActiveKnowledgeBase(tx *gorm.DB, tenantID uint64, id string) error {
|
||||
if tenantID == 0 || id == "" {
|
||||
return ErrKnowledgeBaseNotFound
|
||||
}
|
||||
var kb types.KnowledgeBase
|
||||
err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Select("id").
|
||||
Where("id = ? AND tenant_id = ?", id, tenantID).First(&kb).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrKnowledgeBaseNotFound
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *knowledgeBaseRepository) deleteKnowledgeBaseCategories(ctx context.Context, id string, requireDeleted bool) error {
|
||||
tenantID, ok := types.TenantIDFromContext(ctx)
|
||||
if !ok || tenantID == 0 || id == "" {
|
||||
return errors.New("knowledge base deletion requires tenant and knowledge base IDs")
|
||||
}
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var kb types.KnowledgeBase
|
||||
err := tx.Unscoped().Clauses(clause.Locking{Strength: "UPDATE"}).
|
||||
Select("id", "tenant_id", "deleted_at").Where("id = ? AND tenant_id = ?", id, tenantID).First(&kb).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrKnowledgeBaseNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if requireDeleted && !kb.DeletedAt.Valid {
|
||||
return errors.New("category cleanup requires a deleted knowledge base")
|
||||
}
|
||||
// Include soft-deleted documents: a retry must remove their old relations too.
|
||||
tags := tx.Model(&types.KnowledgeTag{}).Select("id").Where("tenant_id = ? AND knowledge_base_id = ?", tenantID, id)
|
||||
documents := tx.Unscoped().Model(&types.Knowledge{}).Select("id").Where("tenant_id = ? AND knowledge_base_id = ?", tenantID, id)
|
||||
if err := tx.Where("tag_id IN (?) OR knowledge_id IN (?)", tags, documents).
|
||||
Delete(&types.KnowledgeTagRelation{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Where("tenant_id = ? AND knowledge_base_id = ?", tenantID, id).
|
||||
Delete(&types.KnowledgeTag{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if !kb.DeletedAt.Valid {
|
||||
return tx.Where("id = ? AND tenant_id = ?", id, tenantID).Delete(&types.KnowledgeBase{}).Error
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// CountByVectorStoreID counts active knowledge bases that are bound to the
|
||||
// given vector store within a tenant scope.
|
||||
//
|
||||
// Soft-delete filter is applied automatically by GORM because KnowledgeBase
|
||||
// has a gorm.DeletedAt column — we deliberately do not add an explicit
|
||||
// `deleted_at IS NULL` predicate to keep the single source of truth on the
|
||||
// auto-scope.
|
||||
//
|
||||
// Pass db == nil to use the repository's default db handle; pass a *gorm.DB
|
||||
// bound to a transaction (e.g., from db.Transaction) to share the same
|
||||
// write-lock context as the caller. Query column order matches the
|
||||
// composite index idx_knowledge_bases_tenant_vector_store(tenant_id,
|
||||
// vector_store_id).
|
||||
func (r *knowledgeBaseRepository) CountByVectorStoreID(
|
||||
ctx context.Context, db *gorm.DB, tenantID uint64, storeID string,
|
||||
) (int64, error) {
|
||||
if db == nil {
|
||||
db = r.db
|
||||
}
|
||||
var count int64
|
||||
err := db.WithContext(ctx).
|
||||
Model(&types.KnowledgeBase{}).
|
||||
Where("tenant_id = ? AND vector_store_id = ?", tenantID, storeID).
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// CountByModelID counts active knowledge bases that reference modelID in any
|
||||
// model-binding column (scalar fields or JSON config blobs).
|
||||
func (r *knowledgeBaseRepository) CountByModelID(
|
||||
ctx context.Context, tenantID uint64, modelID string,
|
||||
) (int64, error) {
|
||||
var count int64
|
||||
query := r.db.WithContext(ctx).
|
||||
Model(&types.KnowledgeBase{}).
|
||||
Where("tenant_id = ?", tenantID)
|
||||
query = scopeKnowledgeBasesByModelID(query, modelID)
|
||||
err := query.Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// ListModelUsages returns active knowledge bases that reference modelID. It
|
||||
// selects only the columns needed to identify the resource and its bindings;
|
||||
// GORM's default scope excludes soft-deleted rows.
|
||||
func (r *knowledgeBaseRepository) ListModelUsages(
|
||||
ctx context.Context, tenantID uint64, modelID string,
|
||||
) ([]types.ModelUsageResource, error) {
|
||||
rows := make([]*types.KnowledgeBase, 0)
|
||||
query := r.db.WithContext(ctx).
|
||||
Model(&types.KnowledgeBase{}).
|
||||
Select(
|
||||
"id", "name", "embedding_model_id", "summary_model_id",
|
||||
"image_processing_config", "vlm_config", "asr_config", "wiki_config",
|
||||
).
|
||||
Where("tenant_id = ?", tenantID)
|
||||
query = scopeKnowledgeBasesByModelID(query, modelID)
|
||||
if err := query.Order("name ASC, id ASC").Limit(types.ModelUsageListLimit).Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
usages := make([]types.ModelUsageResource, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
bindings := knowledgeBaseModelUsageBindings(row, modelID)
|
||||
if len(bindings) == 0 {
|
||||
continue
|
||||
}
|
||||
usages = append(usages, types.ModelUsageResource{
|
||||
ID: row.ID,
|
||||
Name: row.Name,
|
||||
Bindings: bindings,
|
||||
})
|
||||
}
|
||||
return usages, nil
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupKBDeletePostgres(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
bin := os.Getenv("POSTGRES_TEST_BIN")
|
||||
if bin == "" {
|
||||
bin = "/opt/homebrew/opt/postgresql@16/bin"
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(bin, "initdb")); err != nil {
|
||||
t.Skip("isolated PostgreSQL binaries unavailable; set POSTGRES_TEST_BIN to run lock regression")
|
||||
}
|
||||
dir, err := os.MkdirTemp("/tmp", "kb-delete-pg-")
|
||||
require.NoError(t, err)
|
||||
data := filepath.Join(dir, "data")
|
||||
running := false
|
||||
run := func(name string, args ...string) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
output, err := exec.CommandContext(ctx, filepath.Join(bin, name), args...).CombinedOutput()
|
||||
require.NoError(t, err, "%s: %s", name, output)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if running {
|
||||
run("pg_ctl", "-D", data, "-m", "fast", "-w", "stop")
|
||||
}
|
||||
assert.NoError(t, os.RemoveAll(dir))
|
||||
})
|
||||
run("initdb", "-D", data, "-U", "fixture", "--auth=trust", "--no-locale", "--encoding=UTF8")
|
||||
run("pg_ctl", "-D", data, "-l", filepath.Join(dir, "postgres.log"), "-o",
|
||||
fmt.Sprintf("-F -c listen_addresses='' -k %s -p 55460", dir), "-w", "start")
|
||||
running = true
|
||||
db, err := gorm.Open(postgres.Open(fmt.Sprintf("host=%s port=55460 user=fixture dbname=postgres sslmode=disable application_name=kb-delete-lock-test", dir)), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { assert.NoError(t, sqlDB.Close()) })
|
||||
require.NoError(t, db.Exec(`
|
||||
CREATE TABLE knowledge_bases (id text PRIMARY KEY, tenant_id bigint NOT NULL, deleted_at timestamptz);
|
||||
CREATE TABLE knowledges (id text PRIMARY KEY, tenant_id bigint NOT NULL, knowledge_base_id text NOT NULL, parse_status text, deleted_at timestamptz);
|
||||
CREATE TABLE knowledge_tags (id text PRIMARY KEY, seq_id bigserial UNIQUE, tenant_id bigint NOT NULL, knowledge_base_id text NOT NULL, name text NOT NULL, color text, sort_order int, created_at timestamptz, updated_at timestamptz);
|
||||
CREATE TABLE knowledge_tag_relations (knowledge_id text NOT NULL, tag_id text NOT NULL, created_at timestamptz, PRIMARY KEY (knowledge_id, tag_id));
|
||||
`).Error)
|
||||
return db
|
||||
}
|
||||
|
||||
func awaitKBDeleteBlockedWriter(t *testing.T, db *gorm.DB) {
|
||||
t.Helper()
|
||||
require.Eventually(t, func() bool {
|
||||
var count int64
|
||||
err := db.Raw("SELECT count(*) FROM pg_stat_activity WHERE application_name = 'kb-delete-lock-test' AND wait_event_type = 'Lock'").Scan(&count).Error
|
||||
return err == nil && count > 0
|
||||
}, 5*time.Second, 10*time.Millisecond, "concurrent operation must wait on the KB row lock")
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseDeletePostgresSerializesCategoryWriters(t *testing.T) {
|
||||
db := setupKBDeletePostgres(t)
|
||||
ctx, cancel := context.WithTimeout(kbDeleteContext(1), 30*time.Second)
|
||||
defer cancel()
|
||||
tagRepo := &knowledgeTagRepository{db: db}
|
||||
knowledgeRepo := &knowledgeRepository{db: db}
|
||||
for _, method := range []string{"create", "update", "set relations", "add relations"} {
|
||||
t.Run(method, func(t *testing.T) {
|
||||
require.NoError(t, db.Exec("TRUNCATE knowledge_tag_relations, knowledge_tags, knowledges, knowledge_bases").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id, tenant_id) VALUES ('kb', 1)").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id, tenant_id, knowledge_base_id, parse_status) VALUES ('doc', 1, 'kb', 'completed')").Error)
|
||||
tag := &types.KnowledgeTag{ID: "tag", TenantID: 1, KnowledgeBaseID: "kb", Name: "category"}
|
||||
require.NoError(t, tagRepo.Create(ctx, tag))
|
||||
require.NoError(t, knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{"tag"}))
|
||||
tx := db.WithContext(ctx).Begin()
|
||||
require.NoError(t, tx.Error)
|
||||
defer tx.Rollback()
|
||||
require.NoError(t, (&knowledgeBaseRepository{db: tx}).DeleteKnowledgeBase(ctx, "kb"))
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
switch method {
|
||||
case "create":
|
||||
done <- tagRepo.Create(ctx, &types.KnowledgeTag{ID: "late", TenantID: 1, KnowledgeBaseID: "kb", Name: "late"})
|
||||
case "update":
|
||||
done <- tagRepo.Update(ctx, tag)
|
||||
case "set relations":
|
||||
done <- knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{"tag"})
|
||||
case "add relations":
|
||||
done <- knowledgeRepo.AddKnowledgeTagRelations(ctx, 1, "kb", "doc", []string{"tag"})
|
||||
}
|
||||
}()
|
||||
awaitKBDeleteBlockedWriter(t, db)
|
||||
require.NoError(t, tx.Commit().Error)
|
||||
require.ErrorIs(t, <-done, ErrKnowledgeBaseNotFound)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tag_relations", 0)
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("delete waits for category writer then removes its category", func(t *testing.T) {
|
||||
require.NoError(t, db.Exec("TRUNCATE knowledge_tag_relations, knowledge_tags, knowledges, knowledge_bases").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id, tenant_id) VALUES ('kb', 1)").Error)
|
||||
tx := db.WithContext(ctx).Begin()
|
||||
require.NoError(t, tx.Error)
|
||||
defer tx.Rollback()
|
||||
require.NoError(t, (&knowledgeTagRepository{db: tx}).Create(ctx,
|
||||
&types.KnowledgeTag{ID: "early", TenantID: 1, KnowledgeBaseID: "kb", Name: "early"}))
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- (&knowledgeBaseRepository{db: db}).DeleteKnowledgeBase(ctx, "kb") }()
|
||||
awaitKBDeleteBlockedWriter(t, db)
|
||||
require.NoError(t, tx.Commit().Error)
|
||||
require.NoError(t, <-done)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupKBDeleteTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db := setupKBTestDB(t)
|
||||
require.NoError(t, db.Exec(knowledgeTagTestDDL).Error)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
func kbDeleteContext(tenantID uint64) context.Context {
|
||||
return context.WithValue(context.Background(), types.TenantIDContextKey, tenantID)
|
||||
}
|
||||
|
||||
func seedKBDeleteTag(t *testing.T, db *gorm.DB, kb *types.KnowledgeBase, tagID string) {
|
||||
t.Helper()
|
||||
require.NoError(t, db.Create(&types.KnowledgeTag{
|
||||
ID: tagID, TenantID: kb.TenantID, KnowledgeBaseID: kb.ID, Name: tagID,
|
||||
}).Error)
|
||||
}
|
||||
|
||||
func assertKBDeleteTableCount(t *testing.T, db *gorm.DB, table string, want int64) {
|
||||
t.Helper()
|
||||
var got int64
|
||||
require.NoError(t, db.Table(table).Count(&got).Error)
|
||||
assert.Equal(t, want, got, table)
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseDeleteCleansEmptyCategoriesAndIsIdempotent(t *testing.T) {
|
||||
db := setupKBDeleteTestDB(t)
|
||||
kb := makeKB(nil)
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
seedKBDeleteTag(t, db, kb, "unused-category")
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
ctx := kbDeleteContext(kb.TenantID)
|
||||
|
||||
require.NoError(t, repo.DeleteKnowledgeBase(ctx, kb.ID))
|
||||
var deleted types.KnowledgeBase
|
||||
require.NoError(t, db.Unscoped().First(&deleted, "id = ?", kb.ID).Error)
|
||||
require.True(t, deleted.DeletedAt.Valid)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
|
||||
require.NoError(t, repo.DeleteKnowledgeBase(ctx, kb.ID))
|
||||
require.NoError(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, kb.ID))
|
||||
var retried types.KnowledgeBase
|
||||
require.NoError(t, db.Unscoped().First(&retried, "id = ?", kb.ID).Error)
|
||||
assert.Equal(t, deleted.DeletedAt, retried.DeletedAt)
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseDeleteCleansRelationsAndPreservesOtherResources(t *testing.T) {
|
||||
db := setupKBDeleteTestDB(t)
|
||||
target, sibling, foreign := makeKB(nil), makeKB(nil), makeKB(nil)
|
||||
foreign.TenantID = 2
|
||||
for i, kb := range []*types.KnowledgeBase{target, sibling, foreign} {
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
seedKBDeleteTag(t, db, kb, []string{"target-tag", "sibling-tag", "foreign-tag"}[i])
|
||||
}
|
||||
for _, doc := range []struct {
|
||||
id string
|
||||
kb *types.KnowledgeBase
|
||||
}{
|
||||
{"target-live", target}, {"target-deleted", target}, {"sibling-doc", sibling}, {"foreign-doc", foreign},
|
||||
} {
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id, tenant_id, knowledge_base_id) VALUES (?, ?, ?)",
|
||||
doc.id, doc.kb.TenantID, doc.kb.ID).Error)
|
||||
}
|
||||
require.NoError(t, db.Exec("UPDATE knowledges SET deleted_at = ? WHERE id = ?", time.Now(), "target-deleted").Error)
|
||||
for _, rel := range []types.KnowledgeTagRelation{
|
||||
{KnowledgeID: "target-live", TagID: "target-tag"},
|
||||
{KnowledgeID: "target-deleted", TagID: "target-tag"},
|
||||
// Repair partial old cleanup even when its category or document is already absent.
|
||||
{KnowledgeID: "target-deleted", TagID: "already-removed-tag"},
|
||||
{KnowledgeID: "already-removed-doc", TagID: "target-tag"},
|
||||
{KnowledgeID: "sibling-doc", TagID: "sibling-tag"},
|
||||
{KnowledgeID: "foreign-doc", TagID: "foreign-tag"},
|
||||
} {
|
||||
require.NoError(t, db.Create(&rel).Error)
|
||||
}
|
||||
require.NoError(t, db.Exec("INSERT INTO chunks (id, tenant_id, knowledge_base_id, tag_id) VALUES (?, ?, ?, ?)",
|
||||
"target-chunk", target.TenantID, target.ID, "target-tag").Error)
|
||||
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
require.NoError(t, repo.DeleteKnowledgeBase(kbDeleteContext(1), target.ID))
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 2)
|
||||
assertKBDeleteTableCount(t, db, "knowledges", 4)
|
||||
assertKBDeleteTableCount(t, db, "chunks", 1)
|
||||
var relations []types.KnowledgeTagRelation
|
||||
require.NoError(t, db.Order("knowledge_id").Find(&relations).Error)
|
||||
require.Len(t, relations, 2)
|
||||
assert.Equal(t, "foreign-doc", relations[0].KnowledgeID)
|
||||
assert.Equal(t, "foreign-tag", relations[0].TagID)
|
||||
assert.Equal(t, "sibling-doc", relations[1].KnowledgeID)
|
||||
assert.Equal(t, "sibling-tag", relations[1].TagID)
|
||||
for _, kb := range []*types.KnowledgeBase{sibling, foreign} {
|
||||
got := reloadKB(t, db, kb.ID)
|
||||
assert.False(t, got.DeletedAt.Valid)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseDeleteRejectsUnscopedOrWrongTenant(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ctx context.Context
|
||||
emptyID bool
|
||||
}{
|
||||
{"missing tenant", context.Background(), false},
|
||||
{"zero tenant", kbDeleteContext(0), false},
|
||||
{"wrong tenant", kbDeleteContext(2), false},
|
||||
{"empty ID", kbDeleteContext(1), true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
db := setupKBDeleteTestDB(t)
|
||||
kb := makeKB(nil)
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
seedKBDeleteTag(t, db, kb, "retained-tag")
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
id := kb.ID
|
||||
if tt.emptyID {
|
||||
id = ""
|
||||
}
|
||||
require.Error(t, repo.DeleteKnowledgeBase(tt.ctx, id))
|
||||
require.Error(t, repo.CleanupDeletedKnowledgeBaseCategories(tt.ctx, id))
|
||||
assert.False(t, reloadKB(t, db, kb.ID).DeletedAt.Valid)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseCategoryCleanupRequiresDeletedKB(t *testing.T) {
|
||||
db := setupKBDeleteTestDB(t)
|
||||
kb := makeKB(nil)
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
seedKBDeleteTag(t, db, kb, "retained-tag")
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
ctx := kbDeleteContext(1)
|
||||
require.ErrorContains(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, kb.ID), "requires a deleted")
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 1)
|
||||
require.NoError(t, db.Delete(kb).Error)
|
||||
require.NoError(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, kb.ID))
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
|
||||
require.ErrorIs(t, repo.CleanupDeletedKnowledgeBaseCategories(ctx, "missing-kb"), ErrKnowledgeBaseNotFound)
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseDeleteRollsBackCategoryAndKBChanges(t *testing.T) {
|
||||
for _, trigger := range []struct{ name, sql string }{
|
||||
{"category failure", "CREATE TRIGGER fail_delete BEFORE DELETE ON knowledge_tags BEGIN SELECT RAISE(ABORT, 'category failed'); END"},
|
||||
{"KB failure", "CREATE TRIGGER fail_delete BEFORE UPDATE OF deleted_at ON knowledge_bases BEGIN SELECT RAISE(ABORT, 'KB failed'); END"},
|
||||
} {
|
||||
t.Run(trigger.name, func(t *testing.T) {
|
||||
db := setupKBDeleteTestDB(t)
|
||||
kb := makeKB(nil)
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
seedKBDeleteTag(t, db, kb, "retained-tag")
|
||||
require.NoError(t, db.Create(&types.KnowledgeTagRelation{KnowledgeID: "doc", TagID: "retained-tag"}).Error)
|
||||
require.NoError(t, db.Exec(trigger.sql).Error)
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
require.Error(t, repo.DeleteKnowledgeBase(kbDeleteContext(1), kb.ID))
|
||||
assert.False(t, reloadKB(t, db, kb.ID).DeletedAt.Valid)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 1)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tag_relations", 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestKnowledgeBaseDeleteRejectsLateCategoryWrites(t *testing.T) {
|
||||
db := setupKBDeleteTestDB(t)
|
||||
kb := makeKB(nil)
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
ctx := kbDeleteContext(1)
|
||||
tagRepo := &knowledgeTagRepository{db: db}
|
||||
knowledgeRepo := &knowledgeRepository{db: db}
|
||||
tag := &types.KnowledgeTag{ID: "category", TenantID: 1, KnowledgeBaseID: kb.ID, Name: "original"}
|
||||
require.NoError(t, tagRepo.Create(ctx, tag))
|
||||
tag.Name = "updated"
|
||||
require.NoError(t, tagRepo.Update(ctx, tag))
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id, tenant_id, knowledge_base_id) VALUES (?, ?, ?)", "doc", 1, kb.ID).Error)
|
||||
require.NoError(t, knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{tag.ID}))
|
||||
kbRepo := &knowledgeBaseRepository{db: db}
|
||||
require.NoError(t, kbRepo.DeleteKnowledgeBase(ctx, kb.ID))
|
||||
|
||||
require.ErrorIs(t, tagRepo.Update(ctx, tag), ErrKnowledgeBaseNotFound)
|
||||
require.ErrorIs(t, tagRepo.Create(ctx, &types.KnowledgeTag{ID: "late", TenantID: 1, KnowledgeBaseID: kb.ID, Name: "late"}), ErrKnowledgeBaseNotFound)
|
||||
require.ErrorIs(t, knowledgeRepo.SetKnowledgeTags(ctx, "doc", []string{tag.ID}), ErrKnowledgeBaseNotFound)
|
||||
require.ErrorIs(t, knowledgeRepo.AddKnowledgeTagRelations(ctx, 1, kb.ID, "doc", []string{tag.ID}), ErrKnowledgeBaseNotFound)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tags", 0)
|
||||
assertKBDeleteTableCount(t, db, "knowledge_tag_relations", 0)
|
||||
}
|
||||
@@ -0,0 +1,312 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// knowledgeBasesTestDDL mirrors the `knowledge_bases` section of
|
||||
// migrations/sqlite/000000_init.up.sql. We inline the DDL here instead of
|
||||
// using GORM AutoMigrate because KnowledgeBase carries fields tagged with
|
||||
// `type:jsonb`, which AutoMigrate does not map cleanly onto SQLite.
|
||||
const knowledgeBasesTestDDL = `
|
||||
CREATE TABLE IF NOT EXISTS knowledge_bases (
|
||||
id VARCHAR(36) PRIMARY KEY,
|
||||
name VARCHAR(255) NOT NULL,
|
||||
description TEXT,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
type VARCHAR(32) NOT NULL DEFAULT 'document',
|
||||
chunking_config TEXT NOT NULL DEFAULT '{}',
|
||||
image_processing_config TEXT NOT NULL DEFAULT '{}',
|
||||
embedding_model_id VARCHAR(64) NOT NULL,
|
||||
summary_model_id VARCHAR(64) NOT NULL,
|
||||
cos_config TEXT NOT NULL DEFAULT '{}',
|
||||
storage_provider_config TEXT DEFAULT NULL,
|
||||
vlm_config TEXT NOT NULL DEFAULT '{}',
|
||||
extract_config TEXT NULL DEFAULT NULL,
|
||||
faq_config TEXT,
|
||||
question_generation_config TEXT NULL,
|
||||
auto_tag_config TEXT NULL,
|
||||
is_temporary BOOLEAN NOT NULL DEFAULT 0,
|
||||
version_state TEXT NOT NULL DEFAULT '',
|
||||
is_pinned INTEGER NOT NULL DEFAULT 0,
|
||||
pinned_at DATETIME NULL,
|
||||
asr_config TEXT,
|
||||
vector_store_id VARCHAR(36),
|
||||
storage_backend_id VARCHAR(36),
|
||||
wiki_config TEXT,
|
||||
indexing_strategy TEXT,
|
||||
creator_id VARCHAR(36),
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME
|
||||
);
|
||||
`
|
||||
|
||||
// setupKBTestDB creates an in-memory SQLite database containing the
|
||||
// knowledge_bases table with the vector_store_id column included, so that
|
||||
// tests can exercise the GORM behavior of the VectorStoreID field.
|
||||
func setupKBTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Exec(knowledgeBasesTestDDL).Error)
|
||||
return db
|
||||
}
|
||||
|
||||
// makeKB builds a minimal KnowledgeBase suitable for insert. Only the columns
|
||||
// required by the schema above and any field we explicitly care about are set;
|
||||
// the rest are left at their Go zero values and serialize to "{}" via their
|
||||
// custom JSON valuers.
|
||||
func makeKB(vectorStoreID *string) *types.KnowledgeBase {
|
||||
return &types.KnowledgeBase{
|
||||
ID: uuid.New().String(),
|
||||
Name: "test-kb",
|
||||
TenantID: 1,
|
||||
EmbeddingModelID: "test-embed",
|
||||
SummaryModelID: "test-summary",
|
||||
VectorStoreID: vectorStoreID,
|
||||
}
|
||||
}
|
||||
|
||||
func kbStrPtr(s string) *string { return &s }
|
||||
|
||||
// reloadKB fetches the KnowledgeBase back from the database using its ID.
|
||||
func reloadKB(t *testing.T, db *gorm.DB, id string) types.KnowledgeBase {
|
||||
t.Helper()
|
||||
var got types.KnowledgeBase
|
||||
require.NoError(t, db.First(&got, "id = ?", id).Error)
|
||||
return got
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_VectorStoreID_Save_Immutable verifies that db.Save(kb)
|
||||
// — which WeKnora's UpdateKnowledgeBase repository path uses today — does
|
||||
// NOT overwrite the vector_store_id column, thanks to the GORM
|
||||
// `<-:create` tag. This is the primary automated defense for the H1 review
|
||||
// finding (silent overwrite of the binding via full-struct UPDATE).
|
||||
func TestKnowledgeBase_VectorStoreID_Save_Immutable(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
|
||||
original := makeKB(kbStrPtr("store-A"))
|
||||
require.NoError(t, db.Create(original).Error)
|
||||
|
||||
// Attempt to overwrite through the canonical Save() path.
|
||||
mutated := *original
|
||||
mutated.Name = "updated-name"
|
||||
mutated.VectorStoreID = kbStrPtr("store-B")
|
||||
require.NoError(t, db.Save(&mutated).Error)
|
||||
|
||||
reloaded := reloadKB(t, db, original.ID)
|
||||
|
||||
// Non-immutable fields must update.
|
||||
assert.Equal(t, "updated-name", reloaded.Name, "Name must update through Save")
|
||||
|
||||
// Immutable field must NOT update.
|
||||
require.NotNil(t, reloaded.VectorStoreID, "VectorStoreID must remain set to the original value")
|
||||
assert.Equal(t, "store-A", *reloaded.VectorStoreID,
|
||||
"db.Save must not overwrite vector_store_id (enforced by GORM `<-:create` tag)")
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_VectorStoreID_Updates_Immutable verifies that
|
||||
// db.Updates(struct{...}) also honors `<-:create` and silently ignores the
|
||||
// vector_store_id column.
|
||||
func TestKnowledgeBase_VectorStoreID_Updates_Immutable(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
|
||||
original := makeKB(kbStrPtr("store-A"))
|
||||
require.NoError(t, db.Create(original).Error)
|
||||
|
||||
// Partial struct update that includes vector_store_id must not take effect.
|
||||
err := db.Model(&types.KnowledgeBase{}).
|
||||
Where("id = ?", original.ID).
|
||||
Updates(types.KnowledgeBase{
|
||||
Name: "updated-name",
|
||||
VectorStoreID: kbStrPtr("store-B"),
|
||||
}).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
reloaded := reloadKB(t, db, original.ID)
|
||||
|
||||
assert.Equal(t, "updated-name", reloaded.Name)
|
||||
require.NotNil(t, reloaded.VectorStoreID)
|
||||
assert.Equal(t, "store-A", *reloaded.VectorStoreID,
|
||||
"db.Updates must not overwrite vector_store_id (enforced by GORM `<-:create` tag)")
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_VectorStoreID_Select_StillImmutable verifies that even
|
||||
// an explicit Select("vector_store_id").Updates(...) does NOT bypass the
|
||||
// `<-:create` tag. GORM's `<-:create` is strict: it withholds write
|
||||
// permission from every UPDATE path, including selective updates. This is
|
||||
// a stronger guarantee than a typical "column is protected unless you
|
||||
// opt in" scheme.
|
||||
//
|
||||
// The implication for Phase 4 (cross-store migration) is that ORM-level
|
||||
// updates cannot move bindings. Migration code must use raw SQL (db.Exec)
|
||||
// — see TestKnowledgeBase_VectorStoreID_RawSQL_EscapeHatch below.
|
||||
func TestKnowledgeBase_VectorStoreID_Select_StillImmutable(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
|
||||
original := makeKB(kbStrPtr("store-A"))
|
||||
require.NoError(t, db.Create(original).Error)
|
||||
|
||||
err := db.Model(&types.KnowledgeBase{}).
|
||||
Where("id = ?", original.ID).
|
||||
Select("vector_store_id").
|
||||
Updates(map[string]interface{}{"vector_store_id": "store-B"}).Error
|
||||
require.NoError(t, err)
|
||||
|
||||
reloaded := reloadKB(t, db, original.ID)
|
||||
require.NotNil(t, reloaded.VectorStoreID)
|
||||
assert.Equal(t, "store-A", *reloaded.VectorStoreID,
|
||||
"Select(\"vector_store_id\").Updates must NOT bypass <-:create; ORM is fully locked")
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_VectorStoreID_RawSQL_EscapeHatch documents that raw
|
||||
// SQL UPDATE bypasses the GORM tag system entirely, as expected. Phase 4
|
||||
// cross-store migration is the only sanctioned caller of this path, and
|
||||
// it must do so through a dedicated, narrowly-scoped repository method
|
||||
// rather than through general Save/Updates/Select flows.
|
||||
func TestKnowledgeBase_VectorStoreID_RawSQL_EscapeHatch(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
|
||||
original := makeKB(kbStrPtr("store-A"))
|
||||
require.NoError(t, db.Create(original).Error)
|
||||
|
||||
require.NoError(t,
|
||||
db.Exec("UPDATE knowledge_bases SET vector_store_id = ? WHERE id = ?",
|
||||
"store-B", original.ID).Error)
|
||||
|
||||
reloaded := reloadKB(t, db, original.ID)
|
||||
require.NotNil(t, reloaded.VectorStoreID)
|
||||
assert.Equal(t, "store-B", *reloaded.VectorStoreID,
|
||||
"raw SQL UPDATE bypasses GORM tags; Phase 4 migration must use this path through a dedicated method")
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_VectorStoreID_Roundtrip_Nil verifies that a nil
|
||||
// VectorStoreID is persisted as SQL NULL and loaded back as nil.
|
||||
func TestKnowledgeBase_VectorStoreID_Roundtrip_Nil(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
|
||||
original := makeKB(nil)
|
||||
require.NoError(t, db.Create(original).Error)
|
||||
|
||||
reloaded := reloadKB(t, db, original.ID)
|
||||
assert.Nil(t, reloaded.VectorStoreID,
|
||||
"nil VectorStoreID must round-trip through the database as nil")
|
||||
|
||||
// Extra assurance: column is actually NULL, not an empty string.
|
||||
var nullCount int64
|
||||
require.NoError(t,
|
||||
db.Table("knowledge_bases").
|
||||
Where("id = ? AND vector_store_id IS NULL", original.ID).
|
||||
Count(&nullCount).Error)
|
||||
assert.Equal(t, int64(1), nullCount, "nil must be stored as SQL NULL")
|
||||
}
|
||||
|
||||
// TestKnowledgeBase_VectorStoreID_Roundtrip_Value verifies that a non-nil
|
||||
// VectorStoreID is persisted and loaded back with the same value.
|
||||
func TestKnowledgeBase_VectorStoreID_Roundtrip_Value(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
|
||||
original := makeKB(kbStrPtr("store-uuid-42"))
|
||||
require.NoError(t, db.Create(original).Error)
|
||||
|
||||
reloaded := reloadKB(t, db, original.ID)
|
||||
require.NotNil(t, reloaded.VectorStoreID)
|
||||
assert.Equal(t, "store-uuid-42", *reloaded.VectorStoreID)
|
||||
}
|
||||
|
||||
// TestCountByVectorStoreID covers the binding-count helper used by the
|
||||
// VectorStore delete guard. Verifies tenant scoping, the GORM auto soft-delete
|
||||
// filter (no explicit `deleted_at IS NULL` literal in the query), and the
|
||||
// edge case where a row stores an empty-string vector_store_id (SQLite
|
||||
// treats "" and NULL differently — we want both excluded from a non-empty
|
||||
// storeID lookup).
|
||||
func TestCountByVectorStoreID(t *testing.T) {
|
||||
db := setupKBTestDB(t)
|
||||
repo := &knowledgeBaseRepository{db: db}
|
||||
ctx := t.Context()
|
||||
|
||||
// Tenant 1: two active rows bound to store-A, one row bound to store-B.
|
||||
kbA1 := makeKB(kbStrPtr("store-A"))
|
||||
kbA2 := makeKB(kbStrPtr("store-A"))
|
||||
kbB := makeKB(kbStrPtr("store-B"))
|
||||
require.NoError(t, db.Create(kbA1).Error)
|
||||
require.NoError(t, db.Create(kbA2).Error)
|
||||
require.NoError(t, db.Create(kbB).Error)
|
||||
|
||||
// Tenant 2: one row bound to store-A (must not count for tenant 1).
|
||||
kbCross := makeKB(kbStrPtr("store-A"))
|
||||
kbCross.TenantID = 2
|
||||
require.NoError(t, db.Create(kbCross).Error)
|
||||
|
||||
// Soft-deleted row bound to store-A (must not count under auto-scope).
|
||||
kbDeleted := makeKB(kbStrPtr("store-A"))
|
||||
require.NoError(t, db.Create(kbDeleted).Error)
|
||||
require.NoError(t, db.Delete(kbDeleted).Error)
|
||||
|
||||
// Row with empty-string vector_store_id (regression — sqlite quirk).
|
||||
kbEmpty := makeKB(kbStrPtr(""))
|
||||
require.NoError(t, db.Create(kbEmpty).Error)
|
||||
|
||||
t.Run("nil db handle uses default", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), count, "tenant 1 has 2 active KBs on store-A")
|
||||
})
|
||||
|
||||
t.Run("explicit db handle works the same", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, db, 1, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), count)
|
||||
})
|
||||
|
||||
t.Run("different store id", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "store-B")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
})
|
||||
|
||||
t.Run("unknown store id returns zero", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "store-XYZ")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), count)
|
||||
})
|
||||
|
||||
t.Run("different tenant", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 2, "store-A")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
})
|
||||
|
||||
t.Run("non-empty lookup does not match empty-string rows", func(t *testing.T) {
|
||||
count, err := repo.CountByVectorStoreID(ctx, nil, 1, "")
|
||||
require.NoError(t, err)
|
||||
// Only the empty-string-vsid row matches "" exactly (SQLite quirk —
|
||||
// "" is not NULL). The non-empty rows do not.
|
||||
assert.Equal(t, int64(1), count, "empty-string lookup matches only the empty row")
|
||||
})
|
||||
|
||||
t.Run("shared tx handle returns same result twice", func(t *testing.T) {
|
||||
err := db.Transaction(func(tx *gorm.DB) error {
|
||||
c1, err := repo.CountByVectorStoreID(ctx, tx, 1, "store-A")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c2, err := repo.CountByVectorStoreID(ctx, tx, 1, "store-A")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
assert.Equal(t, c1, c2, "two reads within the same tx must agree")
|
||||
assert.Equal(t, int64(2), c1)
|
||||
return nil
|
||||
})
|
||||
require.NoError(t, err)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// mcpOAuthRepository implements interfaces.MCPOAuthRepository.
|
||||
type mcpOAuthRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewMCPOAuthRepository creates a new MCP OAuth repository.
|
||||
func NewMCPOAuthRepository(db *gorm.DB) interfaces.MCPOAuthRepository {
|
||||
return &mcpOAuthRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) GetClient(
|
||||
ctx context.Context, tenantID uint64, serviceID string,
|
||||
) (*types.MCPOAuthClient, error) {
|
||||
var client types.MCPOAuthClient
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND service_id = ?", tenantID, serviceID).
|
||||
First(&client).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &client, nil
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) SaveClient(ctx context.Context, client *types.MCPOAuthClient) error {
|
||||
client.UpdatedAt = time.Now()
|
||||
return r.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "tenant_id"}, {Name: "service_id"}},
|
||||
DoUpdates: clause.AssignmentColumns([]string{"client_id", "client_secret", "redirect_uri", "updated_at"}),
|
||||
}).
|
||||
Create(client).Error
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) DeleteClient(ctx context.Context, tenantID uint64, serviceID string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND service_id = ?", tenantID, serviceID).
|
||||
Delete(&types.MCPOAuthClient{}).Error
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) GetToken(
|
||||
ctx context.Context, tenantID uint64, userID, serviceID string,
|
||||
) (*types.MCPOAuthToken, error) {
|
||||
return r.GetTokenForPrincipal(ctx, tenantID, types.Principal{
|
||||
Type: types.PrincipalWebUser,
|
||||
ID: userID,
|
||||
}, serviceID)
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) GetTokenForPrincipal(
|
||||
ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
|
||||
) (*types.MCPOAuthToken, error) {
|
||||
principal = principal.Normalize()
|
||||
if !principal.Valid() {
|
||||
return nil, nil
|
||||
}
|
||||
var token types.MCPOAuthToken
|
||||
err := r.db.WithContext(ctx).
|
||||
Where(
|
||||
"tenant_id = ? AND principal_type = ? AND principal_id = ? AND service_id = ?",
|
||||
tenantID, principal.Type, principal.ID, serviceID,
|
||||
).
|
||||
First(&token).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &token, nil
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) SaveToken(ctx context.Context, token *types.MCPOAuthToken) error {
|
||||
if token.PrincipalType == "" || token.PrincipalID == "" {
|
||||
token.PrincipalType = types.PrincipalWebUser
|
||||
token.PrincipalID = token.UserID
|
||||
}
|
||||
return r.SaveTokenForPrincipal(ctx, token)
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) SaveTokenForPrincipal(ctx context.Context, token *types.MCPOAuthToken) error {
|
||||
if token.PrincipalType == "" || token.PrincipalID == "" {
|
||||
return fmt.Errorf("mcp oauth token requires principal_type and principal_id")
|
||||
}
|
||||
if token.UserID == "" {
|
||||
token.UserID = (types.Principal{Type: token.PrincipalType, ID: token.PrincipalID}).StorageID()
|
||||
}
|
||||
token.UpdatedAt = time.Now()
|
||||
return r.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{
|
||||
{Name: "tenant_id"},
|
||||
{Name: "principal_type"},
|
||||
{Name: "principal_id"},
|
||||
{Name: "service_id"},
|
||||
},
|
||||
DoUpdates: clause.AssignmentColumns([]string{
|
||||
"user_id", "access_token", "refresh_token", "token_type", "expires_at", "updated_at",
|
||||
}),
|
||||
}).
|
||||
Create(token).Error
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) DeleteToken(
|
||||
ctx context.Context, tenantID uint64, userID, serviceID string,
|
||||
) error {
|
||||
return r.DeleteTokenForPrincipal(ctx, tenantID, types.Principal{
|
||||
Type: types.PrincipalWebUser,
|
||||
ID: userID,
|
||||
}, serviceID)
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) DeleteTokenForPrincipal(
|
||||
ctx context.Context, tenantID uint64, principal types.Principal, serviceID string,
|
||||
) error {
|
||||
principal = principal.Normalize()
|
||||
if !principal.Valid() {
|
||||
return nil
|
||||
}
|
||||
return r.db.WithContext(ctx).
|
||||
Where(
|
||||
"tenant_id = ? AND principal_type = ? AND principal_id = ? AND service_id = ?",
|
||||
tenantID, principal.Type, principal.ID, serviceID,
|
||||
).
|
||||
Delete(&types.MCPOAuthToken{}).Error
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) TryAcquireTokenRefreshLease(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
principal types.Principal,
|
||||
serviceID, leaseID string,
|
||||
leaseUntil time.Time,
|
||||
) (bool, error) {
|
||||
principal = principal.Normalize()
|
||||
if !principal.Valid() || leaseID == "" {
|
||||
return false, nil
|
||||
}
|
||||
now := time.Now()
|
||||
result := r.db.WithContext(ctx).
|
||||
Model(&types.MCPOAuthToken{}).
|
||||
Where(
|
||||
"tenant_id = ? AND principal_type = ? AND principal_id = ? AND service_id = ?",
|
||||
tenantID, principal.Type, principal.ID, serviceID,
|
||||
).
|
||||
Where("refresh_lease_until IS NULL OR refresh_lease_until < ?", now).
|
||||
Updates(map[string]interface{}{
|
||||
"refresh_lease_id": leaseID,
|
||||
"refresh_lease_until": leaseUntil,
|
||||
})
|
||||
return result.RowsAffected == 1, result.Error
|
||||
}
|
||||
|
||||
func (r *mcpOAuthRepository) ReleaseTokenRefreshLease(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
principal types.Principal,
|
||||
serviceID, leaseID string,
|
||||
) error {
|
||||
principal = principal.Normalize()
|
||||
if !principal.Valid() || leaseID == "" {
|
||||
return nil
|
||||
}
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.MCPOAuthToken{}).
|
||||
Where(
|
||||
"tenant_id = ? AND principal_type = ? AND principal_id = ? AND service_id = ? AND refresh_lease_id = ?",
|
||||
tenantID, principal.Type, principal.ID, serviceID, leaseID,
|
||||
).
|
||||
Updates(map[string]interface{}{
|
||||
"refresh_lease_id": "",
|
||||
"refresh_lease_until": nil,
|
||||
}).Error
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func newMCPOAuthTestRepo(t *testing.T) *mcpOAuthRepository {
|
||||
t.Helper()
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&types.MCPOAuthClient{}, &types.MCPOAuthToken{}))
|
||||
|
||||
return NewMCPOAuthRepository(db).(*mcpOAuthRepository)
|
||||
}
|
||||
|
||||
func TestMCPOAuthRepositoryTokenForPrincipalIsolated(t *testing.T) {
|
||||
repo := newMCPOAuthTestRepo(t)
|
||||
ctx := context.Background()
|
||||
expiresAt := time.Now().Add(time.Hour).UTC()
|
||||
|
||||
webPrincipal := types.Principal{Type: types.PrincipalWebUser, ID: "u1"}
|
||||
apiPrincipal := types.Principal{Type: types.PrincipalAPIExternalUser, ID: "7:external-u1"}
|
||||
|
||||
require.NoError(t, repo.SaveTokenForPrincipal(ctx, &types.MCPOAuthToken{
|
||||
TenantID: 7,
|
||||
UserID: webPrincipal.StorageID(),
|
||||
PrincipalType: webPrincipal.Type,
|
||||
PrincipalID: webPrincipal.ID,
|
||||
ServiceID: "svc1",
|
||||
AccessToken: "web-token",
|
||||
RefreshToken: "web-refresh",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: expiresAt,
|
||||
}))
|
||||
require.NoError(t, repo.SaveTokenForPrincipal(ctx, &types.MCPOAuthToken{
|
||||
TenantID: 7,
|
||||
UserID: apiPrincipal.StorageID(),
|
||||
PrincipalType: apiPrincipal.Type,
|
||||
PrincipalID: apiPrincipal.ID,
|
||||
ServiceID: "svc1",
|
||||
AccessToken: "api-token",
|
||||
RefreshToken: "api-refresh",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: expiresAt,
|
||||
}))
|
||||
|
||||
webToken, err := repo.GetTokenForPrincipal(ctx, 7, webPrincipal, "svc1")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "web-token", webToken.AccessToken)
|
||||
|
||||
apiToken, err := repo.GetTokenForPrincipal(ctx, 7, apiPrincipal, "svc1")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "api-token", apiToken.AccessToken)
|
||||
}
|
||||
|
||||
func TestMCPOAuthRepositoryLegacyUserTokenUsesWebPrincipal(t *testing.T) {
|
||||
repo := newMCPOAuthTestRepo(t)
|
||||
ctx := context.Background()
|
||||
|
||||
require.NoError(t, repo.SaveToken(ctx, &types.MCPOAuthToken{
|
||||
TenantID: 7,
|
||||
UserID: "u1",
|
||||
ServiceID: "svc1",
|
||||
AccessToken: "legacy-token",
|
||||
RefreshToken: "legacy-refresh",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: time.Now().Add(time.Hour).UTC(),
|
||||
}))
|
||||
|
||||
token, err := repo.GetTokenForPrincipal(ctx, 7, types.Principal{Type: types.PrincipalWebUser, ID: "u1"}, "svc1")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "legacy-token", token.AccessToken)
|
||||
}
|
||||
|
||||
func TestMCPOAuthRepositoryRefreshLeaseHasSingleOwner(t *testing.T) {
|
||||
repo := newMCPOAuthTestRepo(t)
|
||||
ctx := context.Background()
|
||||
principal := types.Principal{Type: types.PrincipalWebUser, ID: "u1"}
|
||||
require.NoError(t, repo.SaveTokenForPrincipal(ctx, &types.MCPOAuthToken{
|
||||
TenantID: 7,
|
||||
UserID: principal.StorageID(),
|
||||
PrincipalType: principal.Type,
|
||||
PrincipalID: principal.ID,
|
||||
ServiceID: "svc1",
|
||||
AccessToken: "access",
|
||||
RefreshToken: "refresh",
|
||||
ExpiresAt: time.Now().Add(-time.Minute),
|
||||
}))
|
||||
|
||||
first, err := repo.TryAcquireTokenRefreshLease(
|
||||
ctx, 7, principal, "svc1", "lease-1", time.Now().Add(time.Minute),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.True(t, first)
|
||||
|
||||
second, err := repo.TryAcquireTokenRefreshLease(
|
||||
ctx, 7, principal, "svc1", "lease-2", time.Now().Add(time.Minute),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.False(t, second)
|
||||
|
||||
// A non-owner cannot release the current owner's lease.
|
||||
require.NoError(t, repo.ReleaseTokenRefreshLease(ctx, 7, principal, "svc1", "lease-2"))
|
||||
second, err = repo.TryAcquireTokenRefreshLease(
|
||||
ctx, 7, principal, "svc1", "lease-2", time.Now().Add(time.Minute),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.False(t, second)
|
||||
|
||||
require.NoError(t, repo.ReleaseTokenRefreshLease(ctx, 7, principal, "svc1", "lease-1"))
|
||||
second, err = repo.TryAcquireTokenRefreshLease(
|
||||
ctx, 7, principal, "svc1", "lease-2", time.Now().Add(time.Minute),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.True(t, second)
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// mcpServiceRepository implements the MCPServiceRepository interface
|
||||
type mcpServiceRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewMCPServiceRepository creates a new MCP service repository
|
||||
func NewMCPServiceRepository(db *gorm.DB) interfaces.MCPServiceRepository {
|
||||
return &mcpServiceRepository{db: db}
|
||||
}
|
||||
|
||||
// Create creates a new MCP service
|
||||
func (r *mcpServiceRepository) Create(ctx context.Context, service *types.MCPService) error {
|
||||
return r.db.WithContext(ctx).Create(service).Error
|
||||
}
|
||||
|
||||
// GetByID retrieves an MCP service by ID and tenant ID
|
||||
// Builtin MCP services are visible to all tenants
|
||||
func (r *mcpServiceRepository) GetByID(ctx context.Context, tenantID uint64, id string) (*types.MCPService, error) {
|
||||
var service types.MCPService
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("id = ?", id).
|
||||
Where("tenant_id = ? OR is_builtin = true", tenantID).
|
||||
First(&service).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &service, nil
|
||||
}
|
||||
|
||||
// List retrieves all MCP services for a tenant
|
||||
// Includes builtin MCP services visible to all tenants
|
||||
func (r *mcpServiceRepository) List(ctx context.Context, tenantID uint64) ([]*types.MCPService, error) {
|
||||
var services []*types.MCPService
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? OR is_builtin = true", tenantID).
|
||||
Order("created_at DESC").
|
||||
Find(&services).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return services, nil
|
||||
}
|
||||
|
||||
// ListEnabled retrieves all enabled MCP services for a tenant
|
||||
// Includes enabled builtin MCP services visible to all tenants
|
||||
func (r *mcpServiceRepository) ListEnabled(ctx context.Context, tenantID uint64) ([]*types.MCPService, error) {
|
||||
var services []*types.MCPService
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("(tenant_id = ? OR is_builtin = true) AND enabled = ?", tenantID, true).
|
||||
Order("created_at DESC").
|
||||
Find(&services).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return services, nil
|
||||
}
|
||||
|
||||
// ListByIDs retrieves MCP services by multiple IDs for a tenant
|
||||
// Includes builtin MCP services visible to all tenants
|
||||
func (r *mcpServiceRepository) ListByIDs(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
ids []string,
|
||||
) ([]*types.MCPService, error) {
|
||||
if len(ids) == 0 {
|
||||
return []*types.MCPService{}, nil
|
||||
}
|
||||
|
||||
var services []*types.MCPService
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("(tenant_id = ? OR is_builtin = true) AND id IN ?", tenantID, ids).
|
||||
Find(&services).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return services, nil
|
||||
}
|
||||
|
||||
// Update updates an MCP service
|
||||
func (r *mcpServiceRepository) Update(ctx context.Context, service *types.MCPService) error {
|
||||
// Build update map with only non-zero fields (except enabled which should always be updated if set)
|
||||
updateMap := make(map[string]interface{})
|
||||
updateMap["updated_at"] = service.UpdatedAt
|
||||
|
||||
// Always include enabled field if it's being updated (service layer ensures it's set correctly)
|
||||
updateMap["enabled"] = service.Enabled
|
||||
|
||||
if service.Name != "" {
|
||||
updateMap["name"] = service.Name
|
||||
}
|
||||
// Description can be empty, so we check if it's different from existing
|
||||
// For now, we'll always update it if provided
|
||||
updateMap["description"] = service.Description
|
||||
|
||||
if service.TransportType != "" {
|
||||
updateMap["transport_type"] = service.TransportType
|
||||
}
|
||||
if service.URL != nil {
|
||||
updateMap["url"] = *service.URL
|
||||
}
|
||||
if service.StdioConfig != nil {
|
||||
updateMap["stdio_config"] = service.StdioConfig
|
||||
}
|
||||
if service.EnvVars != nil {
|
||||
updateMap["env_vars"] = service.EnvVars
|
||||
}
|
||||
if service.Headers != nil {
|
||||
updateMap["headers"] = service.Headers
|
||||
}
|
||||
if service.AuthConfig != nil {
|
||||
updateMap["auth_config"] = service.AuthConfig
|
||||
}
|
||||
if service.AdvancedConfig != nil {
|
||||
updateMap["advanced_config"] = service.AdvancedConfig
|
||||
}
|
||||
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.MCPService{}).
|
||||
Where("id = ? AND tenant_id = ?", service.ID, service.TenantID).
|
||||
Updates(updateMap).Error
|
||||
}
|
||||
|
||||
// Delete deletes an MCP service (soft delete)
|
||||
func (r *mcpServiceRepository) Delete(ctx context.Context, tenantID uint64, id string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("id = ? AND tenant_id = ?", id, tenantID).
|
||||
Delete(&types.MCPService{}).Error
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// MCPToolApprovalRepository implements interfaces.MCPToolApprovalRepository.
|
||||
type MCPToolApprovalRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewMCPToolApprovalRepository creates a repository backed by GORM.
|
||||
func NewMCPToolApprovalRepository(db *gorm.DB) interfaces.MCPToolApprovalRepository {
|
||||
return &MCPToolApprovalRepository{db: db}
|
||||
}
|
||||
|
||||
// ListByService returns all stored approval rows for an MCP service (may be empty).
|
||||
func (r *MCPToolApprovalRepository) ListByService(ctx context.Context, tenantID uint64, serviceID string) ([]*types.MCPToolApproval, error) {
|
||||
var rows []*types.MCPToolApproval
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND service_id = ?", tenantID, serviceID).
|
||||
Order("tool_name ASC").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list mcp tool approvals: %w", err)
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
// IsRequired returns true when a row exists with require_approval = true.
|
||||
func (r *MCPToolApprovalRepository) IsRequired(ctx context.Context, tenantID uint64, serviceID, toolName string) (bool, error) {
|
||||
var row types.MCPToolApproval
|
||||
err := r.db.WithContext(ctx).
|
||||
Select("require_approval").
|
||||
Where("tenant_id = ? AND service_id = ? AND tool_name = ?", tenantID, serviceID, toolName).
|
||||
First(&row).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("get mcp tool approval: %w", err)
|
||||
}
|
||||
return row.RequireApproval, nil
|
||||
}
|
||||
|
||||
// Upsert creates or updates the approval flag for a tool atomically.
|
||||
// Uses ON CONFLICT against the (tenant_id, service_id, tool_name) unique index
|
||||
// so concurrent writers don't race the prior SELECT-then-INSERT path into
|
||||
// duplicate-key 500s.
|
||||
func (r *MCPToolApprovalRepository) Upsert(ctx context.Context, row *types.MCPToolApproval) error {
|
||||
if row == nil {
|
||||
return errors.New("row is nil")
|
||||
}
|
||||
if row.ID == "" {
|
||||
row.ID = uuid.New().String()
|
||||
}
|
||||
now := time.Now()
|
||||
row.UpdatedAt = now
|
||||
if row.CreatedAt.IsZero() {
|
||||
row.CreatedAt = now
|
||||
}
|
||||
err := r.db.WithContext(ctx).Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{
|
||||
{Name: "tenant_id"},
|
||||
{Name: "service_id"},
|
||||
{Name: "tool_name"},
|
||||
},
|
||||
DoUpdates: clause.Assignments(map[string]interface{}{
|
||||
"require_approval": row.RequireApproval,
|
||||
"updated_at": now,
|
||||
}),
|
||||
}).Create(row).Error
|
||||
if err != nil {
|
||||
return fmt.Errorf("upsert mcp tool approval: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,186 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/google/uuid"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func currentExtractionSubject(subject *types.MemorySubject, attempt string) bool {
|
||||
return subject != nil && subject.Enabled && subject.ExtractAttemptID == attempt &&
|
||||
subject.ExtractStartedAt != nil && subject.ExtractLeaseUntil != nil && subject.ExtractLeaseUntil.After(time.Now())
|
||||
}
|
||||
|
||||
// Lock the reservation with each model-output mutation. A recovered task or
|
||||
// subject deletion cannot interleave between this check and its item write.
|
||||
func (r *memoryRepository) withExtractionWrite(ctx context.Context, scope interfaces.MemoryScope, write func(*gorm.DB) error) error {
|
||||
attempt := types.MemoryExtractionAttempt(ctx)
|
||||
if attempt == "" {
|
||||
return write(r.db.WithContext(ctx))
|
||||
}
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var subject types.MemorySubject
|
||||
if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID).
|
||||
Clauses(forUpdateClause()).First(&subject).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if !currentExtractionSubject(&subject, attempt) {
|
||||
return errMemoryAttemptLost
|
||||
}
|
||||
return write(tx)
|
||||
})
|
||||
}
|
||||
|
||||
var errMemoryAttemptLost = errors.New("memory extraction attempt is no longer current")
|
||||
|
||||
func extractionRetryDelay(failures int) time.Duration {
|
||||
// Persist the counter; a process restart or a new chat must not reset it.
|
||||
if failures > 9 {
|
||||
return 6 * time.Hour
|
||||
}
|
||||
if failures < 1 {
|
||||
failures = 1
|
||||
}
|
||||
delay := time.Minute * time.Duration(1<<uint(failures-1))
|
||||
if delay > 6*time.Hour {
|
||||
return 6 * time.Hour
|
||||
}
|
||||
return delay
|
||||
}
|
||||
|
||||
func failedExtractionUpdates(subject *types.MemorySubject, now time.Time) map[string]interface{} {
|
||||
failures := subject.ExtractFailures + 1
|
||||
if failures > types.MaxMemoryExtractionFailures {
|
||||
failures = types.MaxMemoryExtractionFailures
|
||||
}
|
||||
var nextAt *time.Time
|
||||
if failures < types.MaxMemoryExtractionFailures {
|
||||
next := now.Add(extractionRetryDelay(failures))
|
||||
nextAt = &next
|
||||
}
|
||||
return map[string]interface{}{
|
||||
"extract_attempt_id": "", "extract_lease_until": nil, "extract_started_at": nil,
|
||||
"extract_scheduled_at": nil, "extract_failures": failures,
|
||||
"extract_next_at": nextAt, "updated_at": now,
|
||||
}
|
||||
}
|
||||
|
||||
// ReserveExtraction commits the task identity before sending its wake-up to
|
||||
// Redis or Lite. A lost wake-up remains discoverable after the lease expires.
|
||||
func (r *memoryRepository) ReserveExtraction(ctx context.Context, scope interfaces.MemoryScope, payload types.MemoryExtractPayload, delay, lease time.Duration) (*types.MemoryExtractPayload, error) {
|
||||
var reserved *types.MemoryExtractPayload
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var subject types.MemorySubject
|
||||
query := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID)
|
||||
if err := query.Clauses(forUpdateClause()).First(&subject).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
updates := map[string]interface{}{}
|
||||
if payload.ChatModelID != "" {
|
||||
subject.ExtractChatModelID = payload.ChatModelID
|
||||
updates["extract_chat_model_id"] = payload.ChatModelID
|
||||
}
|
||||
if payload.Language != "" {
|
||||
subject.ExtractLanguage = payload.Language
|
||||
updates["extract_language"] = payload.Language
|
||||
}
|
||||
active := subject.ExtractLeaseUntil != nil && subject.ExtractLeaseUntil.After(now)
|
||||
backingOff := subject.ExtractNextAt != nil && subject.ExtractNextAt.After(now)
|
||||
if subject.Enabled && subject.ExtractFailures < types.MaxMemoryExtractionFailures && len(subject.PendingSessions) > 0 && !active && !backingOff {
|
||||
if subject.ExtractAttemptID != "" {
|
||||
for key, value := range failedExtractionUpdates(&subject, now) {
|
||||
updates[key] = value
|
||||
}
|
||||
} else {
|
||||
payload.AttemptID = uuid.NewString()
|
||||
payload.TenantID, payload.SubjectID = scope.TenantID, scope.SubjectID
|
||||
payload.SessionID = subject.PendingSessions[0]
|
||||
payload.ChatModelID, payload.Language = subject.ExtractChatModelID, subject.ExtractLanguage
|
||||
updates["extract_attempt_id"] = payload.AttemptID
|
||||
updates["extract_scheduled_at"] = now
|
||||
updates["extract_started_at"] = nil
|
||||
updates["extract_lease_until"] = now.Add(delay + lease)
|
||||
updates["extract_next_at"] = now.Add(delay)
|
||||
reserved = &payload
|
||||
}
|
||||
}
|
||||
if len(updates) == 0 {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&types.MemorySubject{}).Where("id = ?", subject.ID).Updates(updates).Error
|
||||
})
|
||||
return reserved, err
|
||||
}
|
||||
|
||||
func (r *memoryRepository) StartExtraction(ctx context.Context, scope interfaces.MemoryScope, attemptID string, lease time.Duration) (bool, error) {
|
||||
if attemptID == "" {
|
||||
return false, nil
|
||||
}
|
||||
now := time.Now()
|
||||
result := r.scoped(ctx, scope).Model(&types.MemorySubject{}).
|
||||
Where("extract_failures < ?", types.MaxMemoryExtractionFailures).
|
||||
Where("enabled = ? AND extract_attempt_id = ? AND extract_started_at IS NULL AND extract_lease_until > ?", true, attemptID, now).
|
||||
Updates(map[string]interface{}{"extract_started_at": now, "extract_lease_until": now.Add(lease), "updated_at": now})
|
||||
return result.RowsAffected == 1, result.Error
|
||||
}
|
||||
|
||||
func (r *memoryRepository) FailExtraction(ctx context.Context, scope interfaces.MemoryScope, attemptID string) error {
|
||||
if attemptID == "" {
|
||||
return nil
|
||||
}
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var subject types.MemorySubject
|
||||
if err := tx.Where("tenant_id = ? AND subject_id = ?", scope.TenantID, scope.SubjectID).
|
||||
Clauses(forUpdateClause()).First(&subject).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if subject.ExtractAttemptID != attemptID {
|
||||
return nil
|
||||
}
|
||||
return tx.Model(&types.MemorySubject{}).Where("id = ?", subject.ID).
|
||||
Updates(failedExtractionUpdates(&subject, time.Now())).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *memoryRepository) ListDueExtractions(ctx context.Context, limit int) ([]*types.MemorySubject, error) {
|
||||
if limit < 1 || limit > 32 {
|
||||
limit = 32
|
||||
}
|
||||
now := time.Now()
|
||||
var subjects []*types.MemorySubject
|
||||
err := r.db.WithContext(ctx).Where("enabled = ? AND pending_sessions IS NOT NULL AND pending_sessions <> ?", true, "[]").
|
||||
Where("extract_failures < ?", types.MaxMemoryExtractionFailures).
|
||||
Where("extract_next_at IS NULL OR extract_next_at <= ?", now).
|
||||
Where("extract_lease_until IS NULL OR extract_lease_until <= ?", now).
|
||||
// Give tasks queued by the previous binary time to finish during rollout.
|
||||
Where("extract_attempt_id <> '' OR extract_scheduled_at IS NULL OR extract_scheduled_at <= ?", now.Add(-10*time.Minute)).
|
||||
Order("extract_next_at ASC NULLS FIRST, updated_at ASC, id ASC").Limit(limit).Find(&subjects).Error
|
||||
return subjects, err
|
||||
}
|
||||
|
||||
func (r *memoryRepository) DeferExtraction(ctx context.Context, scope interfaces.MemoryScope, delay time.Duration) error {
|
||||
return r.scoped(ctx, scope).Model(&types.MemorySubject{}).
|
||||
Where("extract_attempt_id = ''").
|
||||
Updates(map[string]interface{}{"extract_next_at": time.Now().Add(delay), "extract_scheduled_at": nil}).Error
|
||||
}
|
||||
|
||||
func (r *memoryRepository) PauseExtraction(ctx context.Context, scope interfaces.MemoryScope, attemptID string, delay time.Duration) error {
|
||||
return r.scoped(ctx, scope).Model(&types.MemorySubject{}).Where("extract_attempt_id = ?", attemptID).
|
||||
Updates(map[string]interface{}{
|
||||
"extract_attempt_id": "", "extract_lease_until": nil, "extract_started_at": nil,
|
||||
"extract_scheduled_at": nil, "extract_next_at": time.Now().Add(delay), "updated_at": time.Now(),
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (r *memoryRepository) ResetExhaustedExtraction(ctx context.Context, scope interfaces.MemoryScope) (bool, error) {
|
||||
result := r.scoped(ctx, scope).Model(&types.MemorySubject{}).
|
||||
Where("enabled = ? AND extract_failures >= ? AND extract_attempt_id = ''", true, types.MaxMemoryExtractionFailures).
|
||||
Updates(map[string]interface{}{"extract_failures": 0, "extract_next_at": nil, "extract_scheduled_at": nil, "updated_at": time.Now()})
|
||||
return result.RowsAffected == 1, result.Error
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestMemoryRecoveryReservationsAndExecutionAreExclusive(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
require.NoError(t, db.AutoMigrate(&types.MemorySubject{}))
|
||||
repo := NewMemoryRepository(db)
|
||||
ctx := context.Background()
|
||||
scope := interfaces.MemoryScope{TenantID: 1, SubjectID: "web_user:alice"}
|
||||
_, err = repo.EnsureSubject(ctx, scope)
|
||||
require.NoError(t, err)
|
||||
_, _, err = repo.EnqueuePendingSession(ctx, scope, "s", time.Minute)
|
||||
require.NoError(t, err)
|
||||
var wg sync.WaitGroup
|
||||
results := make(chan *types.MemoryExtractPayload, 12)
|
||||
errors := make(chan error, 12)
|
||||
for range 12 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
payload, err := repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
|
||||
results <- payload
|
||||
errors <- err
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
close(results)
|
||||
close(errors)
|
||||
for err := range errors {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
var winner *types.MemoryExtractPayload
|
||||
for payload := range results {
|
||||
if payload != nil {
|
||||
require.Nil(t, winner)
|
||||
winner = payload
|
||||
}
|
||||
}
|
||||
require.NotNil(t, winner)
|
||||
for i := range 2 {
|
||||
started, err := repo.StartExtraction(ctx, scope, winner.AttemptID, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, i == 0, started)
|
||||
}
|
||||
for failures := 1; failures <= types.MaxMemoryExtractionFailures; failures++ {
|
||||
require.NoError(t, repo.FailExtraction(ctx, scope, winner.AttemptID))
|
||||
subject, err := repo.GetSubject(ctx, scope)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, failures, subject.ExtractFailures)
|
||||
if failures == types.MaxMemoryExtractionFailures {
|
||||
require.Nil(t, subject.ExtractNextAt)
|
||||
winner, err = repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, winner)
|
||||
break
|
||||
}
|
||||
require.InDelta(t, extractionRetryDelay(failures).Seconds(), time.Until(*subject.ExtractNextAt).Seconds(), 1)
|
||||
require.LessOrEqual(t, time.Until(*subject.ExtractNextAt), 6*time.Hour)
|
||||
require.NoError(t, repo.DeferExtraction(ctx, scope, -time.Second))
|
||||
winner, err = repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, winner)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMemoryRecoveryScanIsBoundedAndExcludesIdleScopes(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&types.MemorySubject{}))
|
||||
for i := range 50 {
|
||||
require.NoError(t, db.Create(&types.MemorySubject{ID: fmt.Sprint(i), TenantID: 1, SubjectID: fmt.Sprintf("web_user:%d", i), PendingSessions: types.MemoryPendingSessions{"s"}}).Error)
|
||||
}
|
||||
require.NoError(t, db.Create(&types.MemorySubject{ID: "idle", TenantID: 1, SubjectID: "web_user:idle"}).Error)
|
||||
repo := NewMemoryRepository(db)
|
||||
due, err := repo.ListDueExtractions(context.Background(), 1000)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, due, 32)
|
||||
for _, subject := range due {
|
||||
require.NotEmpty(t, subject.PendingSessions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSQLiteMemoryRecoveryMigrationRoundTrip(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Exec("CREATE TABLE memory_subjects (id TEXT PRIMARY KEY, pending_sessions JSON)").Error)
|
||||
for _, direction := range []string{"up", "down", "up"} {
|
||||
migration, err := os.ReadFile("../../../migrations/sqlite/000217_memory_extraction_recovery." + direction + ".sql")
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Exec(string(migration)).Error)
|
||||
}
|
||||
require.True(t, db.Migrator().HasColumn("memory_subjects", "extract_failures"))
|
||||
require.True(t, db.Migrator().HasColumn("memory_subjects", "session_cursors"))
|
||||
}
|
||||
|
||||
func TestMemoryRecoveryItemWritesRequireLiveReservation(t *testing.T) {
|
||||
for _, change := range []string{"superseded", "expired", "disabled", "deleted", "not started"} {
|
||||
t.Run(change, func(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.AutoMigrate(&types.MemorySubject{}, &types.MemoryItem{}))
|
||||
repo := NewMemoryRepository(db)
|
||||
ctx := context.Background()
|
||||
scope := interfaces.MemoryScope{TenantID: 1, SubjectID: "web_user:alice"}
|
||||
_, err = repo.EnsureSubject(ctx, scope)
|
||||
require.NoError(t, err)
|
||||
_, _, err = repo.EnqueuePendingSession(ctx, scope, "s", time.Minute)
|
||||
require.NoError(t, err)
|
||||
payload, err := repo.ReserveExtraction(ctx, scope, types.MemoryExtractPayload{}, 0, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, payload)
|
||||
started, err := repo.StartExtraction(ctx, scope, payload.AttemptID, time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.True(t, started)
|
||||
workerCtx := types.WithMemoryExtractionAttempt(ctx, payload.AttemptID)
|
||||
original := &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID, Kind: types.MemoryKindFact, Content: "original"}
|
||||
require.NoError(t, repo.CreateItem(workerCtx, original))
|
||||
switch change {
|
||||
case "superseded":
|
||||
err = db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_attempt_id", "new-attempt").Error
|
||||
case "expired":
|
||||
err = db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_lease_until", time.Now().Add(-time.Second)).Error
|
||||
case "disabled":
|
||||
err = repo.UpdateSubjectEnabled(ctx, scope, false)
|
||||
case "deleted":
|
||||
err = db.Where("tenant_id = ?", 1).Delete(&types.MemorySubject{}).Error
|
||||
case "not started":
|
||||
err = db.Model(&types.MemorySubject{}).Where("tenant_id = ?", 1).Update("extract_started_at", nil).Error
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.Error(t, repo.CreateItem(workerCtx, &types.MemoryItem{TenantID: 1, SubjectID: scope.SubjectID, Content: "stale"}))
|
||||
require.Error(t, repo.SupersedeItem(workerCtx, scope, original.ID, ""))
|
||||
_, err = repo.EnsureSubject(workerCtx, scope)
|
||||
require.Error(t, err)
|
||||
var count int64
|
||||
require.NoError(t, db.Model(&types.MemoryItem{}).Count(&count).Error)
|
||||
require.EqualValues(t, 1, count)
|
||||
unchanged, err := repo.GetItem(ctx, scope, original.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, types.MemoryStatusActive, unchanged.Status)
|
||||
if change == "deleted" {
|
||||
subject, err := repo.GetSubject(ctx, scope)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, subject)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,480 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// messageRepository implements the message repository interface
|
||||
type messageRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewMessageRepository creates a new message repository
|
||||
func NewMessageRepository(db *gorm.DB) interfaces.MessageRepository {
|
||||
return &messageRepository{
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
// CreateMessage creates a new message
|
||||
func (r *messageRepository) CreateMessage(
|
||||
ctx context.Context, message *types.Message,
|
||||
) (*types.Message, error) {
|
||||
if err := r.db.WithContext(ctx).Create(message).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return message, nil
|
||||
}
|
||||
|
||||
// GetMessage retrieves a message
|
||||
func (r *messageRepository) GetMessage(
|
||||
ctx context.Context, sessionID string, messageID string,
|
||||
) (*types.Message, error) {
|
||||
var message types.Message
|
||||
if err := r.db.WithContext(ctx).Where(
|
||||
"id = ? AND session_id = ?", messageID, sessionID,
|
||||
).First(&message).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &message, nil
|
||||
}
|
||||
|
||||
// GetMessagesBySession retrieves all messages for a session with pagination
|
||||
func (r *messageRepository) GetMessagesBySession(
|
||||
ctx context.Context, sessionID string, page int, pageSize int,
|
||||
) ([]*types.Message, error) {
|
||||
var messages []*types.Message
|
||||
if err := r.db.WithContext(ctx).Where("session_id = ?", sessionID).Order("created_at ASC").
|
||||
Offset((page - 1) * pageSize).Limit(pageSize).Find(&messages).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// GetRecentMessagesBySession retrieves recent messages for a session
|
||||
func (r *messageRepository) GetRecentMessagesBySession(
|
||||
ctx context.Context, sessionID string, limit int,
|
||||
) ([]*types.Message, error) {
|
||||
var messages []*types.Message
|
||||
if err := r.db.WithContext(ctx).Where(
|
||||
"session_id = ?", sessionID,
|
||||
).Order("created_at DESC").Limit(limit).Find(&messages).Error; err != nil {
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
slices.SortFunc(messages, func(a, b *types.Message) int {
|
||||
cmp := a.CreatedAt.Compare(b.CreatedAt)
|
||||
if cmp == 0 {
|
||||
if a.Role == "user" { // User messages come first
|
||||
return -1
|
||||
}
|
||||
return 1 // Assistant messages come last
|
||||
}
|
||||
return cmp
|
||||
})
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// GetMessagesBySessionBeforeTime retrieves messages from a session created before a specific time
|
||||
func (r *messageRepository) GetMessagesBySessionBeforeTime(
|
||||
ctx context.Context, sessionID string, beforeTime time.Time, limit int,
|
||||
) ([]*types.Message, error) {
|
||||
var messages []*types.Message
|
||||
if err := r.db.WithContext(ctx).Where(
|
||||
"session_id = ? AND created_at < ?", sessionID, beforeTime,
|
||||
).Order("created_at DESC").Limit(limit).Find(&messages).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
slices.SortFunc(messages, func(a, b *types.Message) int {
|
||||
cmp := a.CreatedAt.Compare(b.CreatedAt)
|
||||
if cmp == 0 {
|
||||
if a.Role == "user" { // User messages come first
|
||||
return -1
|
||||
}
|
||||
return 1 // Assistant messages come last
|
||||
}
|
||||
return cmp
|
||||
})
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// ListMessagesBySessionAfterTime returns the oldest messages created after
|
||||
// afterTime, so a caller holding a watermark can walk a session forward
|
||||
// without skipping anything when it has more new messages than one page.
|
||||
func (r *messageRepository) ListMessagesBySessionAfterTime(
|
||||
ctx context.Context, sessionID string, afterTime time.Time, limit int,
|
||||
) ([]*types.Message, error) {
|
||||
var messages []*types.Message
|
||||
query := r.db.WithContext(ctx).Where("session_id = ?", sessionID)
|
||||
if !afterTime.IsZero() {
|
||||
query = query.Where("created_at > ?", afterTime)
|
||||
}
|
||||
if err := query.Order("created_at ASC, id ASC").Limit(limit).Find(&messages).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Memory extraction checkpoints by timestamp. Include the rest of the
|
||||
// boundary timestamp so advancing that checkpoint cannot skip tied rows.
|
||||
if limit > 0 && len(messages) == limit {
|
||||
last := messages[len(messages)-1]
|
||||
var tied []*types.Message
|
||||
if err := r.db.WithContext(ctx).Where("session_id = ? AND created_at = ? AND id > ?", sessionID, last.CreatedAt, last.ID).
|
||||
Order("id ASC").Find(&tied).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
messages = append(messages, tied...)
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
func (r *messageRepository) MemoryOwnedSessionIDs(
|
||||
ctx context.Context, scope interfaces.MemoryScope, sessionIDs []string,
|
||||
) (map[string]bool, error) {
|
||||
owned := map[string]bool{}
|
||||
if !scope.Valid() || len(sessionIDs) == 0 {
|
||||
return owned, nil
|
||||
}
|
||||
query := r.db.WithContext(ctx).Table("sessions").Select("sessions.id").
|
||||
Where("sessions.tenant_id = ? AND sessions.deleted_at IS NULL AND sessions.id IN ?", scope.TenantID, sessionIDs)
|
||||
parts := strings.SplitN(scope.SubjectID, ":", 2)
|
||||
if len(parts) != 2 || parts[1] == "" {
|
||||
return owned, nil
|
||||
}
|
||||
switch parts[0] {
|
||||
case types.PrincipalWebUser:
|
||||
query = query.Where("sessions.user_id = ?", parts[1]).Where(`EXISTS (
|
||||
SELECT 1 FROM users AS memory_user
|
||||
WHERE memory_user.id = sessions.user_id AND memory_user.is_active = ?
|
||||
AND memory_user.deleted_at IS NULL
|
||||
AND (memory_user.is_system_admin = ? OR EXISTS (
|
||||
SELECT 1 FROM tenant_members AS memory_member
|
||||
WHERE memory_member.user_id = memory_user.id AND memory_member.tenant_id = sessions.tenant_id
|
||||
AND memory_member.status = ? AND memory_member.deleted_at IS NULL)))`,
|
||||
true, true, types.TenantMemberStatusActive)
|
||||
case types.PrincipalAPIExternalUser, types.PrincipalEmbedSession:
|
||||
query = query.Where("sessions.user_id = ?", scope.SubjectID)
|
||||
case types.PrincipalIMUser:
|
||||
identity := strings.SplitN(parts[1], ":", 4)
|
||||
if len(identity) != 4 || identity[0] != strconv.FormatUint(scope.TenantID, 10) {
|
||||
return owned, nil
|
||||
}
|
||||
// Shared IM threads do not retain the author on individual messages.
|
||||
// They cannot safely become one participant's personal memory.
|
||||
query = query.Where(`EXISTS (SELECT 1 FROM im_channel_sessions AS im
|
||||
WHERE im.session_id = sessions.id AND im.tenant_id = ? AND im.im_channel_id = ?
|
||||
AND im.platform = ? AND im.user_id = ? AND im.thread_id = '' AND im.deleted_at IS NULL)`,
|
||||
scope.TenantID, identity[1], identity[2], identity[3])
|
||||
default:
|
||||
return owned, nil
|
||||
}
|
||||
var ids []string
|
||||
if err := query.Pluck("sessions.id", &ids).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, id := range ids {
|
||||
owned[id] = true
|
||||
}
|
||||
return owned, nil
|
||||
}
|
||||
|
||||
func (r *messageRepository) MemoryTurnStart(
|
||||
ctx context.Context, scope interfaces.MemoryScope, sessionID, messageID string,
|
||||
) (time.Time, error) {
|
||||
owned, err := r.MemoryOwnedSessionIDs(ctx, scope, []string{sessionID})
|
||||
if err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
if !owned[sessionID] {
|
||||
return time.Time{}, gorm.ErrRecordNotFound
|
||||
}
|
||||
message, err := r.GetMessage(ctx, sessionID, messageID)
|
||||
if err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
if message.Role == "user" {
|
||||
return message.CreatedAt, nil
|
||||
}
|
||||
var input types.Message
|
||||
query := r.db.WithContext(ctx).Where("session_id = ? AND role = ?", sessionID, "user")
|
||||
if message.RequestID != "" {
|
||||
query = query.Where("request_id = ?", message.RequestID)
|
||||
} else {
|
||||
query = query.Where("created_at <= ?", message.CreatedAt)
|
||||
}
|
||||
if err := query.Order("created_at DESC, id DESC").First(&input).Error; err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
return input.CreatedAt, nil
|
||||
}
|
||||
|
||||
// UpdateMessage updates an existing message
|
||||
func (r *messageRepository) UpdateMessage(ctx context.Context, message *types.Message) error {
|
||||
return r.db.WithContext(ctx).Model(&types.Message{}).Where(
|
||||
"id = ? AND session_id = ?", message.ID, message.SessionID,
|
||||
).Updates(message).Error
|
||||
}
|
||||
|
||||
// DeleteMessage deletes a message
|
||||
func (r *messageRepository) DeleteMessage(ctx context.Context, sessionID string, messageID string) error {
|
||||
return r.db.WithContext(ctx).Where(
|
||||
"id = ? AND session_id = ?", messageID, sessionID,
|
||||
).Delete(&types.Message{}).Error
|
||||
}
|
||||
|
||||
// GetFirstMessageOfUser retrieves the first message from a user in a session
|
||||
func (r *messageRepository) GetFirstMessageOfUser(ctx context.Context, sessionID string) (*types.Message, error) {
|
||||
var message types.Message
|
||||
if err := r.db.WithContext(ctx).Where(
|
||||
"session_id = ? and role = ?", sessionID, "user",
|
||||
).Order("created_at ASC").First(&message).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &message, nil
|
||||
}
|
||||
|
||||
// GetMessageByRequestID retrieves a message by request ID
|
||||
func (r *messageRepository) GetMessageByRequestID(
|
||||
ctx context.Context, sessionID string, requestID string,
|
||||
) (*types.Message, error) {
|
||||
var message types.Message
|
||||
|
||||
result := r.db.WithContext(ctx).
|
||||
Where("session_id = ? AND request_id = ?", sessionID, requestID).
|
||||
First(&message)
|
||||
|
||||
if result.Error != nil {
|
||||
if result.Error == gorm.ErrRecordNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, result.Error
|
||||
}
|
||||
|
||||
return &message, nil
|
||||
}
|
||||
|
||||
// SearchMessagesByKeyword searches messages by keyword (ILIKE) across sessions for a tenant
|
||||
func (r *messageRepository) SearchMessagesByKeyword(
|
||||
ctx context.Context, tenantID uint64, ownerID, keyword string, sessionIDs []string, limit int,
|
||||
) ([]*types.MessageWithSession, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
|
||||
var results []*types.MessageWithSession
|
||||
|
||||
query := r.db.WithContext(ctx).
|
||||
Table("messages").
|
||||
Select("messages.*, sessions.title as session_title").
|
||||
Joins("INNER JOIN sessions ON sessions.id = messages.session_id AND sessions.deleted_at IS NULL").
|
||||
Where("sessions.tenant_id = ?", tenantID).
|
||||
Where("messages.deleted_at IS NULL").
|
||||
Where("messages.content ILIKE ?", "%"+escapeLikeKeyword(keyword)+"%")
|
||||
|
||||
// Matches the scoping used when listing sessions, including the legacy
|
||||
// allowance for tenant-level sessions created before per-user ownership.
|
||||
if ownerID != "" {
|
||||
query = query.Where("(sessions.user_id = ? OR sessions.user_id IS NULL OR sessions.user_id = '')", ownerID)
|
||||
}
|
||||
|
||||
if len(sessionIDs) > 0 {
|
||||
query = query.Where("messages.session_id IN ?", sessionIDs)
|
||||
}
|
||||
|
||||
if err := query.Order("messages.created_at DESC").Limit(limit).Find(&results).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// OwnedSessionIDs narrows a set of session ids to the ones this person owns.
|
||||
//
|
||||
// The vector path finds messages through a shared knowledge base that has no
|
||||
// notion of who wrote them, so ownership has to be re-established here before
|
||||
// anything is returned.
|
||||
func (r *messageRepository) OwnedSessionIDs(
|
||||
ctx context.Context, tenantID uint64, ownerID string, sessionIDs []string,
|
||||
) (map[string]bool, error) {
|
||||
owned := make(map[string]bool, len(sessionIDs))
|
||||
if len(sessionIDs) == 0 {
|
||||
return owned, nil
|
||||
}
|
||||
var ids []string
|
||||
query := r.db.WithContext(ctx).
|
||||
Table("sessions").
|
||||
Select("id").
|
||||
Where("tenant_id = ?", tenantID).
|
||||
Where("deleted_at IS NULL").
|
||||
Where("id IN ?", sessionIDs)
|
||||
if ownerID != "" {
|
||||
query = query.Where("(user_id = ? OR user_id IS NULL OR user_id = '')", ownerID)
|
||||
}
|
||||
if err := query.Pluck("id", &ids).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, id := range ids {
|
||||
owned[id] = true
|
||||
}
|
||||
return owned, nil
|
||||
}
|
||||
|
||||
// GetMessagesByKnowledgeIDs retrieves messages by their associated Knowledge IDs
|
||||
func (r *messageRepository) GetMessagesByKnowledgeIDs(
|
||||
ctx context.Context, knowledgeIDs []string,
|
||||
) ([]*types.MessageWithSession, error) {
|
||||
if len(knowledgeIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var results []*types.MessageWithSession
|
||||
if err := r.db.WithContext(ctx).
|
||||
Table("messages").
|
||||
Select("messages.*, sessions.title as session_title").
|
||||
Joins("INNER JOIN sessions ON sessions.id = messages.session_id AND sessions.deleted_at IS NULL").
|
||||
Where("messages.deleted_at IS NULL").
|
||||
Where("messages.knowledge_id IN ?", knowledgeIDs).
|
||||
Find(&results).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// GetMessagesByRequestIDs retrieves messages by their request IDs (used to fetch Q&A pair partners)
|
||||
func (r *messageRepository) GetMessagesByRequestIDs(
|
||||
ctx context.Context, requestIDs []string,
|
||||
) ([]*types.MessageWithSession, error) {
|
||||
if len(requestIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
var results []*types.MessageWithSession
|
||||
if err := r.db.WithContext(ctx).
|
||||
Table("messages").
|
||||
Select("messages.*, sessions.title as session_title").
|
||||
Joins("INNER JOIN sessions ON sessions.id = messages.session_id AND sessions.deleted_at IS NULL").
|
||||
Where("messages.deleted_at IS NULL").
|
||||
Where("messages.request_id IN ?", requestIDs).
|
||||
Find(&results).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// GetKnowledgeIDsBySessionID retrieves all knowledge IDs for messages in a session
|
||||
func (r *messageRepository) GetKnowledgeIDsBySessionID(
|
||||
ctx context.Context, sessionID string,
|
||||
) ([]string, error) {
|
||||
var knowledgeIDs []string
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&types.Message{}).
|
||||
Where("session_id = ? AND knowledge_id != '' AND knowledge_id IS NOT NULL AND deleted_at IS NULL", sessionID).
|
||||
Pluck("knowledge_id", &knowledgeIDs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return knowledgeIDs, nil
|
||||
}
|
||||
|
||||
// UpdateMessageImages updates only the images JSONB column for a message.
|
||||
// Uses Select to force GORM to include the column even when struct-based
|
||||
// Updates would otherwise skip custom Valuer types.
|
||||
func (r *messageRepository) UpdateMessageImages(ctx context.Context, sessionID, messageID string, images types.MessageImages) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.Message{}).
|
||||
Where("id = ? AND session_id = ?", messageID, sessionID).
|
||||
Update("images", images).Error
|
||||
}
|
||||
|
||||
// UpdateMessageRenderedContent updates only the rendered_content column for a message.
|
||||
func (r *messageRepository) UpdateMessageRenderedContent(ctx context.Context, sessionID, messageID string, renderedContent string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.Message{}).
|
||||
Where("id = ? AND session_id = ?", messageID, sessionID).
|
||||
Update("rendered_content", renderedContent).Error
|
||||
}
|
||||
|
||||
// DeleteMessagesBySessionID deletes all messages belonging to a session (soft delete)
|
||||
func (r *messageRepository) DeleteMessagesBySessionID(ctx context.Context, sessionID string) error {
|
||||
return r.db.WithContext(ctx).Where("session_id = ?", sessionID).Delete(&types.Message{}).Error
|
||||
}
|
||||
|
||||
// UpdateMessageKnowledgeID updates the knowledge_id field for a message
|
||||
func (r *messageRepository) UpdateMessageKnowledgeID(
|
||||
ctx context.Context, messageID string, knowledgeID string,
|
||||
) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.Message{}).
|
||||
Where("id = ?", messageID).
|
||||
Update("knowledge_id", knowledgeID).Error
|
||||
}
|
||||
|
||||
// GetSessionArtifacts returns every skill-produced MessageArtifact recorded
|
||||
// against any assistant message of the session, in creation order.
|
||||
//
|
||||
// Projection is scoped to the artifacts JSONB column plus created_at (used
|
||||
// to order the flattened output). Assistant messages without artifacts (the
|
||||
// common case) contribute an empty slice and cost nothing extra.
|
||||
func (r *messageRepository) GetSessionArtifacts(
|
||||
ctx context.Context, sessionID string,
|
||||
) (types.MessageArtifacts, error) {
|
||||
if sessionID == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var rows []struct {
|
||||
Artifacts types.MessageArtifacts `gorm:"column:artifacts"`
|
||||
CreatedAt time.Time `gorm:"column:created_at"`
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&types.Message{}).
|
||||
Select("artifacts", "created_at").
|
||||
Where("session_id = ? AND deleted_at IS NULL", sessionID).
|
||||
Order("created_at ASC").
|
||||
Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(rows) == 0 {
|
||||
return types.MessageArtifacts{}, nil
|
||||
}
|
||||
result := make(types.MessageArtifacts, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if len(row.Artifacts) == 0 {
|
||||
continue
|
||||
}
|
||||
result = append(result, row.Artifacts...)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// GetSessionAttachments returns every user-uploaded attachment in creation
|
||||
// order while projecting only the attachments JSON column.
|
||||
func (r *messageRepository) GetSessionAttachments(
|
||||
ctx context.Context, sessionID string,
|
||||
) (types.MessageAttachments, error) {
|
||||
if sessionID == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var rows []struct {
|
||||
Attachments types.MessageAttachments `gorm:"column:attachments"`
|
||||
CreatedAt time.Time `gorm:"column:created_at"`
|
||||
}
|
||||
if err := r.db.WithContext(ctx).
|
||||
Model(&types.Message{}).
|
||||
Select("attachments", "created_at").
|
||||
Where("session_id = ? AND deleted_at IS NULL", sessionID).
|
||||
Order("created_at ASC").
|
||||
Find(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result := make(types.MessageAttachments, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
result = append(result, row.Attachments...)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func addMemoryWebOwner(t *testing.T, db *gorm.DB) {
|
||||
t.Helper()
|
||||
for _, query := range []string{
|
||||
`CREATE TABLE users (id TEXT PRIMARY KEY, is_active BOOLEAN, is_system_admin BOOLEAN, deleted_at DATETIME)`,
|
||||
`CREATE TABLE tenant_members (user_id TEXT, tenant_id INTEGER, status TEXT, deleted_at DATETIME)`,
|
||||
`INSERT INTO users (id, is_active, is_system_admin) VALUES ('alice', true, false)`,
|
||||
`INSERT INTO tenant_members (user_id, tenant_id, status) VALUES ('alice', 7, 'active')`,
|
||||
} {
|
||||
require.NoError(t, db.Exec(query).Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMemoryMessagePageIncludesTimestampBoundary(t *testing.T) {
|
||||
db, sessions := newOwnerScopeDB(t, t.Name())
|
||||
repo := NewMessageRepository(db)
|
||||
base := time.Now().Add(-time.Hour).UTC().Truncate(time.Microsecond)
|
||||
for i := 0; i < 50; i++ {
|
||||
at := base
|
||||
if i == 49 {
|
||||
at = base.Add(time.Second)
|
||||
}
|
||||
require.NoError(t, db.Create(&types.Message{ID: fmt.Sprintf("message-%02d", i), SessionID: sessions["alice"], Role: "user", Content: "test", CreatedAt: at}).Error)
|
||||
}
|
||||
require.NoError(t, db.Create(&types.Message{ID: "other-session", SessionID: sessions["bob"], Role: "user", Content: "other", CreatedAt: base}).Error)
|
||||
first, err := repo.ListMessagesBySessionAfterTime(context.Background(), sessions["alice"], time.Time{}, 41)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, first, 49)
|
||||
for _, row := range first {
|
||||
require.Equal(t, sessions["alice"], row.SessionID)
|
||||
}
|
||||
next, err := repo.ListMessagesBySessionAfterTime(context.Background(), sessions["alice"], base, 41)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, next, 1)
|
||||
require.True(t, base.Add(time.Second).Equal(next[0].CreatedAt))
|
||||
}
|
||||
|
||||
func TestMemorySessionOwnershipAndCompletedTurnAttribution(t *testing.T) {
|
||||
db, ids := newOwnerScopeDB(t, t.Name())
|
||||
addMemoryWebOwner(t, db)
|
||||
// Web sessions persist the raw account ID, while memory uses a typed ID.
|
||||
require.NoError(t, db.Model(&types.Session{}).Where("id IN ?", []string{ids["alice"], ids["other"]}).Update("user_id", "alice").Error)
|
||||
repo := NewMessageRepository(db)
|
||||
scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
|
||||
ctx := context.Background()
|
||||
owned, err := repo.MemoryOwnedSessionIDs(ctx, scope, []string{ids["alice"], ids["bob"], ids["legacy"], ids["other"]})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[string]bool{ids["alice"]: true}, owned)
|
||||
base := time.Now().UTC().Add(-time.Hour).Truncate(time.Microsecond)
|
||||
var answerID string
|
||||
for _, row := range []*types.Message{
|
||||
{ID: "input", SessionID: ids["alice"], Role: "user", RequestID: "slow", Content: "slow input", CreatedAt: base},
|
||||
{ID: "unrelated", SessionID: ids["alice"], Role: "user", RequestID: "fast", Content: "later input", CreatedAt: base.Add(time.Second)},
|
||||
{ID: "answer", SessionID: ids["alice"], Role: "assistant", RequestID: "slow", Content: "answer", CreatedAt: base.Add(time.Minute)},
|
||||
} {
|
||||
require.NoError(t, db.Create(row).Error)
|
||||
if row.Role == "assistant" {
|
||||
answerID = row.ID
|
||||
}
|
||||
}
|
||||
start, err := repo.MemoryTurnStart(ctx, scope, ids["alice"], answerID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, start.Equal(base), "use the matching request input, not a later unrelated turn")
|
||||
_, err = repo.MemoryTurnStart(ctx, interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:bob"}, ids["alice"], answerID)
|
||||
require.Error(t, err)
|
||||
require.NoError(t, db.Delete(&types.Session{}, "id = ?", ids["alice"]).Error)
|
||||
owned, err = repo.MemoryOwnedSessionIDs(ctx, scope, []string{ids["alice"]})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, owned)
|
||||
}
|
||||
|
||||
func TestMemoryWebOwnerRequiresCurrentAccountAndWorkspaceAccess(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
change string
|
||||
allowed bool
|
||||
}{
|
||||
{"active member", "SELECT 1", true},
|
||||
{"inactive account", "UPDATE users SET is_active = false", false},
|
||||
{"deleted account", "UPDATE users SET deleted_at = CURRENT_TIMESTAMP", false},
|
||||
{"missing account", "DELETE FROM users", false},
|
||||
{"suspended membership", "UPDATE tenant_members SET status = 'suspended'", false},
|
||||
{"invited membership", "UPDATE tenant_members SET status = 'invited'", false},
|
||||
{"deleted membership", "UPDATE tenant_members SET deleted_at = CURRENT_TIMESTAMP", false},
|
||||
{"missing membership", "DELETE FROM tenant_members", false},
|
||||
{"wrong workspace", "UPDATE tenant_members SET tenant_id = 8", false},
|
||||
{"system admin", "UPDATE users SET is_system_admin = true", true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
db, ids := newOwnerScopeDB(t, t.Name())
|
||||
addMemoryWebOwner(t, db)
|
||||
require.NoError(t, db.Model(&types.Session{}).Where("id = ?", ids["alice"]).Update("user_id", "alice").Error)
|
||||
require.NoError(t, db.Exec(tc.change).Error)
|
||||
if tc.name == "system admin" {
|
||||
require.NoError(t, db.Exec("DELETE FROM tenant_members").Error)
|
||||
}
|
||||
repo := NewMessageRepository(db)
|
||||
scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
|
||||
owned, err := repo.MemoryOwnedSessionIDs(context.Background(), scope, []string{ids["alice"]})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.allowed, owned[ids["alice"]])
|
||||
if tc.name == "system admin" {
|
||||
require.NoError(t, db.Exec("UPDATE users SET is_active = false").Error)
|
||||
owned, err = repo.MemoryOwnedSessionIDs(context.Background(), scope, []string{ids["alice"]})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, owned)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMemoryNonWebOwnerDoesNotRequireWebAccount(t *testing.T) {
|
||||
for _, kind := range []string{types.PrincipalAPIExternalUser, types.PrincipalEmbedSession} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
db, ids := newOwnerScopeDB(t, t.Name())
|
||||
scope := interfaces.MemoryScope{TenantID: 7, SubjectID: kind + ":external"}
|
||||
require.NoError(t, db.Model(&types.Session{}).Where("id = ?", ids["alice"]).Update("user_id", scope.SubjectID).Error)
|
||||
owned, err := NewMessageRepository(db).MemoryOwnedSessionIDs(context.Background(), scope, []string{ids["alice"]})
|
||||
require.NoError(t, err)
|
||||
require.True(t, owned[ids["alice"]])
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// MemoryAllowedMessageIDs uses per-turn Agent identity; a session can switch
|
||||
// Agents and its last-request settings cannot authorize older messages.
|
||||
func (r *messageRepository) MemoryAllowedMessageIDs(ctx context.Context, scope interfaces.MemoryScope, sessionID string, messages []*types.Message) (map[string]bool, error) {
|
||||
allowed := map[string]bool{}
|
||||
owned, err := r.MemoryOwnedSessionIDs(ctx, scope, []string{sessionID})
|
||||
if err != nil || !owned[sessionID] {
|
||||
return allowed, err
|
||||
}
|
||||
for _, input := range messages {
|
||||
if input == nil || input.Role != "user" || input.SessionID != sessionID {
|
||||
continue
|
||||
}
|
||||
var answer types.Message
|
||||
query := r.db.WithContext(ctx).Where("session_id = ? AND role = ?", sessionID, "assistant")
|
||||
if input.RequestID != "" {
|
||||
query = query.Where("request_id = ?", input.RequestID)
|
||||
} else {
|
||||
// Without a request identity, only an adjacent assistant can safely
|
||||
// attribute a historical input. Never borrow another turn's Agent.
|
||||
query = r.db.WithContext(ctx).Where("session_id = ? AND created_at > ?", sessionID, input.CreatedAt)
|
||||
}
|
||||
if err := query.Order("created_at ASC, id ASC").First(&answer).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
continue
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if answer.Role != "assistant" {
|
||||
continue
|
||||
}
|
||||
if !answer.IsCompleted {
|
||||
// Keep the checkpoint behind a turn still streaming in this session.
|
||||
return nil, errors.New("memory extraction is waiting for a completed turn")
|
||||
}
|
||||
if answer.AgentID == "" {
|
||||
// Historical ordinary-assistant turns have no custom Agent binding.
|
||||
allowed[input.ID] = true
|
||||
continue
|
||||
}
|
||||
agentTenant := answer.AgentTenantID
|
||||
if agentTenant == 0 {
|
||||
agentTenant = scope.TenantID
|
||||
}
|
||||
var agent types.CustomAgent
|
||||
err := r.db.WithContext(ctx).Where("tenant_id = ? AND id = ?", agentTenant, answer.AgentID).First(&agent).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
// Built-in defaults may have no tenant override. A soft-deleted
|
||||
// override must not be recreated through the built-in registry.
|
||||
var deleted int64
|
||||
if err := r.db.WithContext(ctx).Unscoped().Model(&types.CustomAgent{}).
|
||||
Where("tenant_id = ? AND id = ?", agentTenant, answer.AgentID).Count(&deleted).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
builtin := types.GetBuiltinAgent(answer.AgentID, agentTenant)
|
||||
if deleted != 0 || builtin == nil || agentTenant != scope.TenantID {
|
||||
continue
|
||||
}
|
||||
agent = *builtin
|
||||
} else if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !agent.Enabled || (agent.Config.MemoryEnabled != nil && !*agent.Config.MemoryEnabled) {
|
||||
continue
|
||||
}
|
||||
if agentTenant != scope.TenantID {
|
||||
var shares int64
|
||||
err := r.db.WithContext(ctx).Table("agent_shares").
|
||||
Joins("JOIN organization_tenant_members otm ON otm.organization_id = agent_shares.organization_id").
|
||||
Joins("JOIN organizations ON organizations.id = agent_shares.organization_id AND organizations.deleted_at IS NULL").
|
||||
Where("agent_shares.agent_id = ? AND agent_shares.source_tenant_id = ? AND agent_shares.deleted_at IS NULL AND otm.tenant_id = ?", answer.AgentID, agentTenant, scope.TenantID).
|
||||
Count(&shares).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if shares == 0 {
|
||||
continue
|
||||
}
|
||||
}
|
||||
allowed[input.ID] = true
|
||||
}
|
||||
return allowed, nil
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMemoryAgentPolicyUsesPerTurnIdentityAndRejectsDeletedAgents(t *testing.T) {
|
||||
db, sessions := newOwnerScopeDB(t, t.Name())
|
||||
addMemoryWebOwner(t, db)
|
||||
require.NoError(t, db.AutoMigrate(&types.CustomAgent{}))
|
||||
require.NoError(t, db.Model(&types.Session{}).Where("id = ?", sessions["alice"]).Update("user_id", "alice").Error)
|
||||
disabled := false
|
||||
agents := []*types.CustomAgent{
|
||||
{ID: "allowed", TenantID: 7, Enabled: true},
|
||||
{ID: "disabled", TenantID: 7, Enabled: false},
|
||||
{ID: "opted-out", TenantID: 7, Enabled: true, Config: types.CustomAgentConfig{MemoryEnabled: &disabled}},
|
||||
{ID: "deleted", TenantID: 7, Enabled: true},
|
||||
{ID: "foreign", TenantID: 8, Enabled: true},
|
||||
}
|
||||
for _, agent := range agents {
|
||||
require.NoError(t, db.Create(agent).Error)
|
||||
}
|
||||
require.NoError(t, db.Model(agents[1]).Update("enabled", false).Error)
|
||||
require.NoError(t, db.Delete(agents[3]).Error)
|
||||
base := time.Now().Add(-time.Hour)
|
||||
var inputs []*types.Message
|
||||
for i, agentID := range []string{"allowed", "disabled", "opted-out", "deleted", "missing", "foreign", ""} {
|
||||
input := &types.Message{ID: "user-" + agentID, SessionID: sessions["alice"], Role: "user", RequestID: "request-" + agentID, CreatedAt: base.Add(time.Duration(i) * time.Minute)}
|
||||
answer := &types.Message{ID: "answer-" + agentID, SessionID: input.SessionID, Role: "assistant", RequestID: input.RequestID, AgentID: agentID, IsCompleted: true, CreatedAt: input.CreatedAt.Add(time.Second)}
|
||||
require.NoError(t, db.Create(input).Error)
|
||||
require.NoError(t, db.Create(answer).Error)
|
||||
inputs = append(inputs, input)
|
||||
}
|
||||
repo := NewMessageRepository(db)
|
||||
scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
|
||||
allowed, err := repo.MemoryAllowedMessageIDs(context.Background(), scope, sessions["alice"], inputs)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[string]bool{inputs[0].ID: true, inputs[6].ID: true}, allowed)
|
||||
require.NoError(t, db.Model(agents[0]).Update("enabled", false).Error)
|
||||
allowed, err = repo.MemoryAllowedMessageIDs(context.Background(), scope, sessions["alice"], inputs)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, map[string]bool{inputs[6].ID: true}, allowed)
|
||||
allowed, err = repo.MemoryAllowedMessageIDs(context.Background(), interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:bob"}, sessions["alice"], inputs)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, allowed)
|
||||
}
|
||||
|
||||
func TestMemorySharedAgentRechecksExactSourceAndCurrentMembership(t *testing.T) {
|
||||
db, sessions := newOwnerScopeDB(t, t.Name())
|
||||
addMemoryWebOwner(t, db)
|
||||
require.NoError(t, db.AutoMigrate(&types.CustomAgent{}))
|
||||
require.NoError(t, db.Model(&types.Session{}).Where("id = ?", sessions["alice"]).Update("user_id", "alice").Error)
|
||||
for _, query := range []string{
|
||||
`CREATE TABLE agent_shares (id TEXT PRIMARY KEY, agent_id TEXT, organization_id TEXT, source_tenant_id INTEGER, deleted_at DATETIME)`,
|
||||
`CREATE TABLE organization_tenant_members (organization_id TEXT, tenant_id INTEGER)`,
|
||||
`CREATE TABLE organizations (id TEXT PRIMARY KEY, deleted_at DATETIME)`,
|
||||
`INSERT INTO organizations(id) VALUES ('org')`,
|
||||
`INSERT INTO organization_tenant_members VALUES ('org', 7)`,
|
||||
`INSERT INTO agent_shares(id, agent_id, organization_id, source_tenant_id) VALUES ('share', 'shared', 'org', 84)`,
|
||||
} {
|
||||
require.NoError(t, db.Exec(query).Error)
|
||||
}
|
||||
require.NoError(t, db.Create(&types.CustomAgent{ID: "shared", TenantID: 84, Enabled: true}).Error)
|
||||
require.NoError(t, db.Create(&types.CustomAgent{ID: "shared", TenantID: 42, Enabled: true}).Error)
|
||||
input := &types.Message{ID: "input", SessionID: sessions["alice"], Role: "user", RequestID: "request", CreatedAt: time.Now().Add(-time.Minute)}
|
||||
answer := &types.Message{ID: "answer", SessionID: input.SessionID, Role: "assistant", RequestID: input.RequestID, AgentID: "shared", AgentTenantID: 84, IsCompleted: true, CreatedAt: time.Now()}
|
||||
require.NoError(t, db.Create(input).Error)
|
||||
require.NoError(t, db.Create(answer).Error)
|
||||
repo := NewMessageRepository(db)
|
||||
scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
|
||||
check := func(want bool) {
|
||||
t.Helper()
|
||||
allowed, err := repo.MemoryAllowedMessageIDs(context.Background(), scope, input.SessionID, []*types.Message{input})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, want, allowed[input.ID])
|
||||
}
|
||||
check(true)
|
||||
require.NoError(t, db.Model(answer).Update("agent_tenant_id", 42).Error)
|
||||
check(false)
|
||||
require.NoError(t, db.Model(answer).Update("agent_tenant_id", 84).Error)
|
||||
require.NoError(t, db.Exec("DELETE FROM organization_tenant_members").Error)
|
||||
check(false)
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Sessions are per-user state and the session list has always been scoped that
|
||||
// way. Conversation search was not, which meant a workspace viewer could type a
|
||||
// colleague's project name into the search box and read their chats. These
|
||||
// tests pin the scoping so that cannot come back.
|
||||
|
||||
// Session.BeforeCreate assigns a fresh UUID, so the ids are read back rather
|
||||
// than assumed.
|
||||
func newOwnerScopeDB(t *testing.T, name string) (*gorm.DB, map[string]string) {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open("file:"+name+"?mode=memory&cache=shared"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf("open sqlite: %v", err)
|
||||
}
|
||||
if err := db.AutoMigrate(&types.Session{}, &types.Message{}); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
sessions := map[string]*types.Session{
|
||||
"alice": {TenantID: 7, UserID: "web_user:alice", Title: "Alice 的会话"},
|
||||
"bob": {TenantID: 7, UserID: "web_user:bob", Title: "Bob 的会话"},
|
||||
"legacy": {TenantID: 7, Title: "API 建的会话"},
|
||||
"other": {TenantID: 8, UserID: "web_user:alice", Title: "别的工作区"},
|
||||
}
|
||||
ids := make(map[string]string, len(sessions))
|
||||
for label, session := range sessions {
|
||||
if err := db.Create(session).Error; err != nil {
|
||||
t.Fatalf("insert session: %v", err)
|
||||
}
|
||||
ids[label] = session.ID
|
||||
}
|
||||
return db, ids
|
||||
}
|
||||
|
||||
func TestOwnedSessionIDsExcludesOtherPeople(t *testing.T) {
|
||||
db, ids := newOwnerScopeDB(t, "owner-scope-owned")
|
||||
repo := NewMessageRepository(db)
|
||||
|
||||
owned, err := repo.OwnedSessionIDs(context.Background(), 7, "web_user:alice",
|
||||
[]string{ids["alice"], ids["bob"], ids["legacy"], ids["other"]})
|
||||
if err != nil {
|
||||
t.Fatalf("owned session ids: %v", err)
|
||||
}
|
||||
|
||||
if !owned[ids["alice"]] {
|
||||
t.Error("alice must be able to search her own conversations")
|
||||
}
|
||||
if owned[ids["bob"]] {
|
||||
t.Error("alice must not be able to reach bob's conversations")
|
||||
}
|
||||
if !owned[ids["legacy"]] {
|
||||
t.Error("tenant-level sessions stay reachable, matching how they are listed")
|
||||
}
|
||||
if owned[ids["other"]] {
|
||||
t.Error("a workspace boundary is not something an owner check may cross")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOwnedSessionIDsWithoutAnOwnerKeepsWorkspaceScope(t *testing.T) {
|
||||
db, ids := newOwnerScopeDB(t, "owner-scope-noowner")
|
||||
repo := NewMessageRepository(db)
|
||||
|
||||
// An empty owner means "no per-user narrowing", which is what admin-console
|
||||
// style callers rely on. It must still not leak across workspaces.
|
||||
owned, err := repo.OwnedSessionIDs(context.Background(), 7, "",
|
||||
[]string{ids["alice"], ids["bob"], ids["other"]})
|
||||
if err != nil {
|
||||
t.Fatalf("owned session ids: %v", err)
|
||||
}
|
||||
if !owned[ids["alice"]] || !owned[ids["bob"]] {
|
||||
t.Error("without an owner filter, every session in the workspace qualifies")
|
||||
}
|
||||
if owned[ids["other"]] {
|
||||
t.Error("tenant scoping is not optional")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type messageSuggestionRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func NewMessageSuggestionRepository(db *gorm.DB) interfaces.MessageSuggestionRepository {
|
||||
return &messageSuggestionRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) GetByCacheKey(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
assistantMessageID string,
|
||||
placement string,
|
||||
configHash string,
|
||||
locale string,
|
||||
) (*types.MessageSuggestionSet, error) {
|
||||
var set types.MessageSuggestionSet
|
||||
err := r.db.WithContext(ctx).
|
||||
Where(
|
||||
"tenant_id = ? AND assistant_message_id = ? AND placement = ? AND config_hash = ? AND locale = ?",
|
||||
tenantID, assistantMessageID, placement, configHash, locale,
|
||||
).
|
||||
First(&set).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &set, nil
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) GetByID(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
sessionID string,
|
||||
id string,
|
||||
) (*types.MessageSuggestionSet, error) {
|
||||
var set types.MessageSuggestionSet
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("id = ? AND tenant_id = ? AND session_id = ?", id, tenantID, sessionID).
|
||||
First(&set).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &set, nil
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) AcquireGeneration(
|
||||
ctx context.Context,
|
||||
candidate *types.MessageSuggestionSet,
|
||||
regenerate bool,
|
||||
) (*types.MessageSuggestionSet, bool, error) {
|
||||
now := time.Now()
|
||||
leaseUntil := now.Add(3 * time.Minute)
|
||||
candidate.Status = types.SuggestionStatusGenerating
|
||||
candidate.LeaseUntil = &leaseUntil
|
||||
candidate.Questions = types.SuggestionItems{}
|
||||
|
||||
result := r.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{DoNothing: true}).
|
||||
Create(candidate)
|
||||
if result.Error != nil {
|
||||
return nil, false, result.Error
|
||||
}
|
||||
if result.RowsAffected == 1 {
|
||||
return candidate, true, nil
|
||||
}
|
||||
|
||||
existing, err := r.GetByCacheKey(
|
||||
ctx,
|
||||
candidate.TenantID,
|
||||
candidate.AssistantMessageID,
|
||||
candidate.Placement,
|
||||
candidate.ConfigHash,
|
||||
candidate.Locale,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if existing.Status == types.SuggestionStatusReady && !regenerate {
|
||||
return existing, false, nil
|
||||
}
|
||||
if existing.Status == types.SuggestionStatusSuppressed && !regenerate {
|
||||
return existing, false, nil
|
||||
}
|
||||
if existing.Status == types.SuggestionStatusGenerating && existing.LeaseUntil != nil && existing.LeaseUntil.After(now) {
|
||||
return existing, false, nil
|
||||
}
|
||||
|
||||
query := r.db.WithContext(ctx).Model(&types.MessageSuggestionSet{}).
|
||||
Where("id = ?", existing.ID).
|
||||
Where("status <> ? OR lease_until IS NULL OR lease_until < ?", types.SuggestionStatusGenerating, now)
|
||||
if existing.Status == types.SuggestionStatusReady && regenerate {
|
||||
query = r.db.WithContext(ctx).Model(&types.MessageSuggestionSet{}).
|
||||
Where("id = ? AND status = ?", existing.ID, types.SuggestionStatusReady)
|
||||
}
|
||||
updates := map[string]interface{}{
|
||||
"status": types.SuggestionStatusGenerating,
|
||||
"lease_until": leaseUntil,
|
||||
"suppression_reason": "",
|
||||
"questions": types.SuggestionItems{},
|
||||
"error_code": "",
|
||||
"generated_at": nil,
|
||||
"updated_at": now,
|
||||
}
|
||||
result = query.Updates(updates)
|
||||
if result.Error != nil {
|
||||
return nil, false, result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
current, getErr := r.GetByCacheKey(
|
||||
ctx,
|
||||
candidate.TenantID,
|
||||
candidate.AssistantMessageID,
|
||||
candidate.Placement,
|
||||
candidate.ConfigHash,
|
||||
candidate.Locale,
|
||||
)
|
||||
return current, false, getErr
|
||||
}
|
||||
existing.Status = types.SuggestionStatusGenerating
|
||||
existing.LeaseUntil = &leaseUntil
|
||||
existing.Questions = types.SuggestionItems{}
|
||||
return existing, true, nil
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) Save(ctx context.Context, set *types.MessageSuggestionSet) error {
|
||||
if set == nil {
|
||||
return errors.New("message suggestion set is nil")
|
||||
}
|
||||
return r.db.WithContext(ctx).Save(set).Error
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) CreateEvent(
|
||||
ctx context.Context,
|
||||
event *types.MessageSuggestionEvent,
|
||||
) error {
|
||||
return r.db.WithContext(ctx).Create(event).Error
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) DeleteByMessageID(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
sessionID string,
|
||||
messageID string,
|
||||
) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND session_id = ? AND assistant_message_id = ?", tenantID, sessionID, messageID).
|
||||
Delete(&types.MessageSuggestionSet{}).Error
|
||||
}
|
||||
|
||||
func (r *messageSuggestionRepository) DeleteBySessionID(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
sessionID string,
|
||||
) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND session_id = ?", tenantID, sessionID).
|
||||
Delete(&types.MessageSuggestionSet{}).Error
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// modelRepository implements the model repository interface
|
||||
type modelRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewModelRepository creates a new model repository
|
||||
func NewModelRepository(db *gorm.DB) interfaces.ModelRepository {
|
||||
return &modelRepository{db: db}
|
||||
}
|
||||
|
||||
// Create creates a new model
|
||||
func (r *modelRepository) Create(ctx context.Context, m *types.Model) error {
|
||||
return r.db.WithContext(ctx).Create(m).Error
|
||||
}
|
||||
|
||||
// GetByID retrieves a model by ID
|
||||
func (r *modelRepository) GetByID(ctx context.Context, tenantID uint64, id string) (*types.Model, error) {
|
||||
var m types.Model
|
||||
if err := r.db.WithContext(ctx).Where("id = ?", id).Where(
|
||||
"(tenant_id = ? OR is_builtin = true)", tenantID,
|
||||
).First(&m).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &m, nil
|
||||
}
|
||||
|
||||
// List lists models with optional filtering
|
||||
func (r *modelRepository) List(
|
||||
ctx context.Context, tenantID uint64, modelType types.ModelType, source types.ModelSource,
|
||||
) ([]*types.Model, error) {
|
||||
var models []*types.Model
|
||||
query := r.db.WithContext(ctx).Where(
|
||||
"(tenant_id = ? OR is_builtin = true)", tenantID,
|
||||
)
|
||||
|
||||
if modelType != "" {
|
||||
query = query.Where("type = ?", modelType)
|
||||
}
|
||||
|
||||
if source != "" {
|
||||
query = query.Where("source = ?", source)
|
||||
}
|
||||
|
||||
if err := query.Find(&models).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return models, nil
|
||||
}
|
||||
|
||||
// Update updates a model
|
||||
func (r *modelRepository) Update(ctx context.Context, m *types.Model) error {
|
||||
// Use Select to explicitly update all fields, including zero values like false
|
||||
return r.db.WithContext(ctx).Debug().Model(&types.Model{}).Where(
|
||||
"id = ? AND tenant_id = ?", m.ID, m.TenantID,
|
||||
).Select("*").Updates(m).Error
|
||||
}
|
||||
|
||||
// Delete deletes a model
|
||||
func (r *modelRepository) Delete(ctx context.Context, tenantID uint64, id string) error {
|
||||
return r.db.WithContext(ctx).Where(
|
||||
"id = ? AND tenant_id = ?", id, tenantID,
|
||||
).Delete(&types.Model{}).Error
|
||||
}
|
||||
|
||||
// ClearDefaultByType clears the default flag for all models of a specific type
|
||||
// This is a batch operation that updates all matching records in one query
|
||||
func (r *modelRepository) ClearDefaultByType(
|
||||
ctx context.Context,
|
||||
tenantID uint,
|
||||
modelType types.ModelType,
|
||||
excludeID string,
|
||||
) error {
|
||||
query := r.db.WithContext(ctx).Model(&types.Model{}).Where(
|
||||
"tenant_id = ? AND type = ? AND is_default = ?", tenantID, modelType, true,
|
||||
)
|
||||
|
||||
// If excludeID is provided, exclude that model from the update
|
||||
if excludeID != "" {
|
||||
query = query.Where("id != ?", excludeID)
|
||||
}
|
||||
|
||||
// Batch update: set is_default to false for all matching records
|
||||
return query.Update("is_default", false).Error
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func knowledgeBaseModelUsageBindings(kb *types.KnowledgeBase, modelID string) []types.ModelUsageBinding {
|
||||
bindings := make([]types.ModelUsageBinding, 0, 6)
|
||||
if kb.EmbeddingModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingEmbeddingModel)
|
||||
}
|
||||
if kb.SummaryModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingSummaryModel)
|
||||
}
|
||||
if kb.ImageProcessingConfig.ModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingImageProcessingModel)
|
||||
}
|
||||
if kb.VLMConfig.ModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingVLMModel)
|
||||
}
|
||||
if kb.ASRConfig.ModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingASRModel)
|
||||
}
|
||||
if kb.WikiConfig != nil && kb.WikiConfig.SynthesisModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingWikiSynthesisModel)
|
||||
}
|
||||
return bindings
|
||||
}
|
||||
|
||||
func customAgentModelUsageBindings(agent *types.CustomAgent, modelID string) []types.ModelUsageBinding {
|
||||
bindings := make([]types.ModelUsageBinding, 0, 6)
|
||||
if agent.Config.ModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingChatModel)
|
||||
}
|
||||
if agent.Config.RerankModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingRerankModel)
|
||||
}
|
||||
if agent.Config.VLMModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingVLMModel)
|
||||
}
|
||||
if agent.Config.ASRModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingASRModel)
|
||||
}
|
||||
if agent.Config.QueryUnderstandModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingQueryUnderstandModel)
|
||||
}
|
||||
if agent.Config.QuestionSuggestions != nil &&
|
||||
agent.Config.QuestionSuggestions.FollowUps.ModelID == modelID {
|
||||
bindings = append(bindings, types.ModelUsageBindingFollowUpModel)
|
||||
}
|
||||
return bindings
|
||||
}
|
||||
|
||||
// scopeKnowledgeBasesByModelID filters knowledge_bases rows that reference
|
||||
// modelID in any model-binding field.
|
||||
func scopeKnowledgeBasesByModelID(db *gorm.DB, modelID string) *gorm.DB {
|
||||
if db.Dialector.Name() == "postgres" {
|
||||
return db.Where(
|
||||
"embedding_model_id = ? OR summary_model_id = ? OR "+
|
||||
"image_processing_config->>'model_id' = ? OR "+
|
||||
"vlm_config->>'model_id' = ? OR "+
|
||||
"asr_config->>'model_id' = ? OR "+
|
||||
"wiki_config->>'synthesis_model_id' = ?",
|
||||
modelID, modelID, modelID, modelID, modelID, modelID,
|
||||
)
|
||||
}
|
||||
return db.Where(
|
||||
"embedding_model_id = ? OR summary_model_id = ? OR "+
|
||||
"json_extract(image_processing_config, '$.model_id') = ? OR "+
|
||||
"json_extract(vlm_config, '$.model_id') = ? OR "+
|
||||
"json_extract(asr_config, '$.model_id') = ? OR "+
|
||||
"json_extract(wiki_config, '$.synthesis_model_id') = ?",
|
||||
modelID, modelID, modelID, modelID, modelID, modelID,
|
||||
)
|
||||
}
|
||||
|
||||
// scopeCustomAgentsByModelID filters custom_agents rows whose config JSON
|
||||
// references modelID in any model-binding field.
|
||||
func scopeCustomAgentsByModelID(db *gorm.DB, modelID string) *gorm.DB {
|
||||
if db.Dialector.Name() == "postgres" {
|
||||
return db.Where(
|
||||
"config->>'model_id' = ? OR config->>'rerank_model_id' = ? OR "+
|
||||
"config->>'vlm_model_id' = ? OR config->>'asr_model_id' = ? OR "+
|
||||
"config->>'query_understand_model_id' = ? OR "+
|
||||
"config->'question_suggestions'->'follow_ups'->>'model_id' = ?",
|
||||
modelID, modelID, modelID, modelID, modelID, modelID,
|
||||
)
|
||||
}
|
||||
return db.Where(
|
||||
"json_extract(config, '$.model_id') = ? OR "+
|
||||
"json_extract(config, '$.rerank_model_id') = ? OR "+
|
||||
"json_extract(config, '$.vlm_model_id') = ? OR "+
|
||||
"json_extract(config, '$.asr_model_id') = ? OR "+
|
||||
"json_extract(config, '$.query_understand_model_id') = ? OR "+
|
||||
"json_extract(config, '$.question_suggestions.follow_ups.model_id') = ?",
|
||||
modelID, modelID, modelID, modelID, modelID, modelID,
|
||||
)
|
||||
}
|
||||
|
||||
func scopeCustomAgentsBySandboxConfigID(db *gorm.DB, configID string) *gorm.DB {
|
||||
if db.Dialector.Name() == "postgres" {
|
||||
return db.Where("config->>'sandbox_config_id' = ?", configID)
|
||||
}
|
||||
return db.Where("json_extract(config, '$.sandbox_config_id') = ?", configID)
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const customAgentsTestDDL = `
|
||||
CREATE TABLE IF NOT EXISTS custom_agents (
|
||||
id VARCHAR(36) NOT NULL,
|
||||
name VARCHAR(255) NOT NULL,
|
||||
description TEXT,
|
||||
avatar VARCHAR(64),
|
||||
is_builtin BOOLEAN NOT NULL DEFAULT 0,
|
||||
tenant_id INTEGER NOT NULL,
|
||||
created_by VARCHAR(36),
|
||||
visibility VARCHAR(16) NOT NULL DEFAULT 'global',
|
||||
enabled BOOLEAN NOT NULL DEFAULT 1,
|
||||
config TEXT NOT NULL DEFAULT '{}',
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
deleted_at DATETIME,
|
||||
PRIMARY KEY (id, tenant_id)
|
||||
);
|
||||
`
|
||||
|
||||
func setupModelUsageTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
db := setupKBTestDB(t)
|
||||
require.NoError(t, db.Exec(customAgentsTestDDL).Error)
|
||||
return db
|
||||
}
|
||||
|
||||
func TestCountByModelID_KnowledgeBase(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db := setupModelUsageTestDB(t)
|
||||
repo := NewKnowledgeBaseRepository(db)
|
||||
modelID := "embed-model-1"
|
||||
|
||||
kb := makeKB(nil)
|
||||
kb.EmbeddingModelID = modelID
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
|
||||
count, err := repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
|
||||
count, err = repo.CountByModelID(ctx, 1, "other-model")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), count)
|
||||
|
||||
kb2 := makeKB(nil)
|
||||
kb2.ID = uuid.New().String()
|
||||
kb2.VLMConfig = types.VLMConfig{Enabled: true, ModelID: modelID}
|
||||
require.NoError(t, db.Create(kb2).Error)
|
||||
|
||||
count, err = repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), count)
|
||||
|
||||
require.NoError(t, db.Delete(kb2).Error)
|
||||
count, err = repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
}
|
||||
|
||||
func TestListModelUsages_KnowledgeBase(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db := setupModelUsageTestDB(t)
|
||||
repo := NewKnowledgeBaseRepository(db)
|
||||
modelID := "all-kb-bindings"
|
||||
|
||||
allBindings := makeKB(nil)
|
||||
allBindings.Name = "Zulu KB"
|
||||
allBindings.EmbeddingModelID = modelID
|
||||
allBindings.SummaryModelID = modelID
|
||||
allBindings.ImageProcessingConfig = types.ImageProcessingConfig{ModelID: modelID}
|
||||
allBindings.VLMConfig = types.VLMConfig{Enabled: true, ModelID: modelID}
|
||||
allBindings.ASRConfig = types.ASRConfig{Enabled: true, ModelID: modelID}
|
||||
allBindings.WikiConfig = &types.WikiConfig{SynthesisModelID: modelID}
|
||||
require.NoError(t, db.Create(allBindings).Error)
|
||||
|
||||
vlmOnly := makeKB(nil)
|
||||
vlmOnly.Name = "Alpha KB"
|
||||
vlmOnly.VLMConfig = types.VLMConfig{Enabled: true, ModelID: modelID}
|
||||
require.NoError(t, db.Create(vlmOnly).Error)
|
||||
|
||||
otherTenant := makeKB(nil)
|
||||
otherTenant.Name = "Other tenant"
|
||||
otherTenant.TenantID = 2
|
||||
otherTenant.EmbeddingModelID = modelID
|
||||
require.NoError(t, db.Create(otherTenant).Error)
|
||||
|
||||
deleted := makeKB(nil)
|
||||
deleted.Name = "Deleted KB"
|
||||
deleted.EmbeddingModelID = modelID
|
||||
require.NoError(t, db.Create(deleted).Error)
|
||||
require.NoError(t, db.Delete(deleted).Error)
|
||||
|
||||
usages, err := repo.ListModelUsages(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, usages, 2)
|
||||
assert.Equal(t, "Alpha KB", usages[0].Name)
|
||||
assert.Equal(t, []types.ModelUsageBinding{types.ModelUsageBindingVLMModel}, usages[0].Bindings)
|
||||
assert.Equal(t, "Zulu KB", usages[1].Name)
|
||||
assert.Equal(t, []types.ModelUsageBinding{
|
||||
types.ModelUsageBindingEmbeddingModel,
|
||||
types.ModelUsageBindingSummaryModel,
|
||||
types.ModelUsageBindingImageProcessingModel,
|
||||
types.ModelUsageBindingVLMModel,
|
||||
types.ModelUsageBindingASRModel,
|
||||
types.ModelUsageBindingWikiSynthesisModel,
|
||||
}, usages[1].Bindings)
|
||||
}
|
||||
|
||||
func TestListModelUsages_KnowledgeBaseRespectsLimit(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db := setupModelUsageTestDB(t)
|
||||
repo := NewKnowledgeBaseRepository(db)
|
||||
modelID := "shared-embed"
|
||||
|
||||
for i := 0; i < types.ModelUsageListLimit+3; i++ {
|
||||
kb := makeKB(nil)
|
||||
kb.Name = fmt.Sprintf("KB %03d", i)
|
||||
kb.EmbeddingModelID = modelID
|
||||
require.NoError(t, db.Create(kb).Error)
|
||||
}
|
||||
|
||||
usages, err := repo.ListModelUsages(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, usages, types.ModelUsageListLimit)
|
||||
assert.Equal(t, "KB 000", usages[0].Name)
|
||||
|
||||
count, err := repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(types.ModelUsageListLimit+3), count)
|
||||
}
|
||||
|
||||
func TestCountByModelID_CustomAgent(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db := setupModelUsageTestDB(t)
|
||||
repo := NewCustomAgentRepository(db)
|
||||
modelID := "chat-model-1"
|
||||
|
||||
agent := &types.CustomAgent{
|
||||
ID: uuid.New().String(),
|
||||
Name: "test-agent",
|
||||
TenantID: 1,
|
||||
Config: types.CustomAgentConfig{
|
||||
ModelID: modelID,
|
||||
},
|
||||
}
|
||||
require.NoError(t, repo.CreateAgent(ctx, agent))
|
||||
|
||||
count, err := repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
|
||||
agent2 := &types.CustomAgent{
|
||||
ID: uuid.New().String(),
|
||||
Name: "rerank-agent",
|
||||
TenantID: 1,
|
||||
Config: types.CustomAgentConfig{
|
||||
RerankModelID: modelID,
|
||||
},
|
||||
}
|
||||
require.NoError(t, repo.CreateAgent(ctx, agent2))
|
||||
|
||||
count, err = repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2), count)
|
||||
|
||||
count, err = repo.CountByModelID(ctx, 2, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), count)
|
||||
|
||||
require.NoError(t, repo.DeleteAgent(ctx, agent2.ID, 1))
|
||||
count, err = repo.CountByModelID(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
}
|
||||
|
||||
func TestListModelUsages_CustomAgent(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db := setupModelUsageTestDB(t)
|
||||
repo := NewCustomAgentRepository(db)
|
||||
modelID := "all-agent-bindings"
|
||||
|
||||
allBindings := &types.CustomAgent{
|
||||
ID: uuid.New().String(),
|
||||
Name: "Zulu agent",
|
||||
TenantID: 1,
|
||||
Config: types.CustomAgentConfig{
|
||||
ModelID: modelID,
|
||||
RerankModelID: modelID,
|
||||
VLMModelID: modelID,
|
||||
ASRModelID: modelID,
|
||||
QueryUnderstandModelID: modelID,
|
||||
QuestionSuggestions: &types.QuestionSuggestionConfig{
|
||||
FollowUps: types.FollowUpSuggestionConfig{ModelID: modelID},
|
||||
},
|
||||
},
|
||||
}
|
||||
require.NoError(t, repo.CreateAgent(ctx, allBindings))
|
||||
|
||||
vlmOnly := &types.CustomAgent{
|
||||
ID: uuid.New().String(),
|
||||
Name: "Alpha agent",
|
||||
TenantID: 1,
|
||||
Config: types.CustomAgentConfig{VLMModelID: modelID},
|
||||
}
|
||||
require.NoError(t, repo.CreateAgent(ctx, vlmOnly))
|
||||
|
||||
otherTenant := &types.CustomAgent{
|
||||
ID: uuid.New().String(), Name: "Other tenant", TenantID: 2,
|
||||
Config: types.CustomAgentConfig{ModelID: modelID},
|
||||
}
|
||||
require.NoError(t, repo.CreateAgent(ctx, otherTenant))
|
||||
|
||||
deleted := &types.CustomAgent{
|
||||
ID: uuid.New().String(), Name: "Deleted agent", TenantID: 1,
|
||||
Config: types.CustomAgentConfig{ModelID: modelID},
|
||||
}
|
||||
require.NoError(t, repo.CreateAgent(ctx, deleted))
|
||||
require.NoError(t, repo.DeleteAgent(ctx, deleted.ID, deleted.TenantID))
|
||||
|
||||
usages, err := repo.ListModelUsages(ctx, 1, modelID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, usages, 2)
|
||||
assert.Equal(t, "Alpha agent", usages[0].Name)
|
||||
assert.Equal(t, []types.ModelUsageBinding{types.ModelUsageBindingVLMModel}, usages[0].Bindings)
|
||||
assert.Equal(t, "Zulu agent", usages[1].Name)
|
||||
assert.Equal(t, []types.ModelUsageBinding{
|
||||
types.ModelUsageBindingChatModel,
|
||||
types.ModelUsageBindingRerankModel,
|
||||
types.ModelUsageBindingVLMModel,
|
||||
types.ModelUsageBindingASRModel,
|
||||
types.ModelUsageBindingQueryUnderstandModel,
|
||||
types.ModelUsageBindingFollowUpModel,
|
||||
}, usages[1].Bindings)
|
||||
}
|
||||
|
||||
func TestCustomAgentSandboxConfigReferences(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db := setupModelUsageTestDB(t)
|
||||
repo := NewCustomAgentRepository(db)
|
||||
configID := "sandbox-cfg-1"
|
||||
|
||||
require.NoError(t, db.Exec(
|
||||
`INSERT INTO custom_agents (id, name, tenant_id, config) VALUES
|
||||
(?, ?, ?, ?),
|
||||
(?, ?, ?, ?),
|
||||
(?, ?, ?, ?)`,
|
||||
uuid.New().String(), "analysis-agent", 1, `{"sandbox_config_id":"sandbox-cfg-1"}`,
|
||||
uuid.New().String(), "other-agent", 1, `{"sandbox_config_id":"sandbox-cfg-2"}`,
|
||||
uuid.New().String(), "other-tenant-agent", 2, `{"sandbox_config_id":"sandbox-cfg-1"}`,
|
||||
).Error)
|
||||
|
||||
count, err := repo.CountBySandboxConfigID(ctx, 1, configID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), count)
|
||||
|
||||
names, err := repo.ListNamesBySandboxConfigID(ctx, 1, configID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"analysis-agent"}, names)
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrOrganizationNotFound = errors.New("organization not found")
|
||||
ErrOrgMemberNotFound = errors.New("organization member not found")
|
||||
ErrOrgMemberAlreadyExists = errors.New("member already exists in organization")
|
||||
ErrInviteCodeNotFound = errors.New("invite code not found")
|
||||
ErrInviteCodeExpired = errors.New("invite code has expired")
|
||||
)
|
||||
|
||||
// organizationRepository implements OrganizationRepository.
|
||||
//
|
||||
// Plan 3 of #1303 lifts membership from per-user to per-tenant (see
|
||||
// migration 000045). All "member" methods on this repo now operate on
|
||||
// the (org, tenant) tuple; the underlying table is
|
||||
// organization_tenant_members.
|
||||
type organizationRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewOrganizationRepository creates a new organization repository
|
||||
func NewOrganizationRepository(db *gorm.DB) interfaces.OrganizationRepository {
|
||||
return &organizationRepository{db: db}
|
||||
}
|
||||
|
||||
// Create creates a new organization
|
||||
func (r *organizationRepository) Create(ctx context.Context, org *types.Organization) error {
|
||||
return r.db.WithContext(ctx).Create(org).Error
|
||||
}
|
||||
|
||||
// GetByID gets an organization by ID
|
||||
func (r *organizationRepository) GetByID(ctx context.Context, id string) (*types.Organization, error) {
|
||||
var org types.Organization
|
||||
if err := r.db.WithContext(ctx).Where("id = ?", id).First(&org).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrOrganizationNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &org, nil
|
||||
}
|
||||
|
||||
// GetByInviteCode gets an organization by invite code (returns ErrInviteCodeExpired if code has expired)
|
||||
func (r *organizationRepository) GetByInviteCode(ctx context.Context, inviteCode string) (*types.Organization, error) {
|
||||
var org types.Organization
|
||||
if err := r.db.WithContext(ctx).Where("invite_code = ?", inviteCode).First(&org).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrInviteCodeNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if org.InviteCodeExpiresAt != nil && org.InviteCodeExpiresAt.Before(time.Now()) {
|
||||
return nil, ErrInviteCodeExpired
|
||||
}
|
||||
return &org, nil
|
||||
}
|
||||
|
||||
// ListByTenantID lists organizations whose tenant participates as a member.
|
||||
func (r *organizationRepository) ListByTenantID(ctx context.Context, tenantID uint64) ([]*types.Organization, error) {
|
||||
var orgs []*types.Organization
|
||||
|
||||
err := r.db.WithContext(ctx).
|
||||
Joins("JOIN organization_tenant_members otm ON otm.organization_id = organizations.id").
|
||||
Where("otm.tenant_id = ?", tenantID).
|
||||
Order("organizations.created_at DESC").
|
||||
Find(&orgs).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return orgs, nil
|
||||
}
|
||||
|
||||
// ListSearchable lists organizations that are searchable (open for discovery), optionally filtered by name/description/ID
|
||||
func (r *organizationRepository) ListSearchable(ctx context.Context, query string, limit int) ([]*types.Organization, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
var orgs []*types.Organization
|
||||
q := r.db.WithContext(ctx).Where("searchable = ?", true)
|
||||
if query != "" {
|
||||
pattern := "%" + query + "%"
|
||||
// 支持按名称、描述或空间 ID 搜索,便于区分同名空间
|
||||
q = q.Where("name ILIKE ? OR description ILIKE ? OR id::text ILIKE ?", pattern, pattern, pattern)
|
||||
}
|
||||
err := q.Order("created_at DESC").Limit(limit).Find(&orgs).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return orgs, nil
|
||||
}
|
||||
|
||||
// Update updates an organization (Select ensures zero values like invite_code_validity_days=0 are persisted)
|
||||
func (r *organizationRepository) Update(ctx context.Context, org *types.Organization) error {
|
||||
return r.db.WithContext(ctx).Model(&types.Organization{}).Where("id = ?", org.ID).
|
||||
Select("name", "description", "avatar", "require_approval", "searchable", "invite_code_validity_days", "member_limit", "updated_at").
|
||||
Updates(org).Error
|
||||
}
|
||||
|
||||
// Delete soft deletes an organization
|
||||
func (r *organizationRepository) Delete(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).Where("id = ?", id).Delete(&types.Organization{}).Error
|
||||
}
|
||||
|
||||
// AddTenantMember inserts a new (org, tenant) membership row. Returns
|
||||
// ErrOrgMemberAlreadyExists if a row already exists for this tuple — the
|
||||
// service layer treats that as a no-op when the caller is just confirming
|
||||
// an idempotent join.
|
||||
func (r *organizationRepository) AddTenantMember(ctx context.Context, member *types.OrganizationTenantMember) error {
|
||||
var count int64
|
||||
r.db.WithContext(ctx).Model(&types.OrganizationTenantMember{}).
|
||||
Where("organization_id = ? AND tenant_id = ?", member.OrganizationID, member.TenantID).
|
||||
Count(&count)
|
||||
|
||||
if count > 0 {
|
||||
return ErrOrgMemberAlreadyExists
|
||||
}
|
||||
|
||||
return r.db.WithContext(ctx).Create(member).Error
|
||||
}
|
||||
|
||||
// RemoveTenantMember removes the (org, tenant) membership row.
|
||||
func (r *organizationRepository) RemoveTenantMember(ctx context.Context, orgID string, tenantID uint64) error {
|
||||
result := r.db.WithContext(ctx).
|
||||
Where("organization_id = ? AND tenant_id = ?", orgID, tenantID).
|
||||
Delete(&types.OrganizationTenantMember{})
|
||||
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return ErrOrgMemberNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateTenantMemberRole updates the role for a (org, tenant) membership.
|
||||
func (r *organizationRepository) UpdateTenantMemberRole(ctx context.Context, orgID string, tenantID uint64, role types.OrgMemberRole) error {
|
||||
result := r.db.WithContext(ctx).
|
||||
Model(&types.OrganizationTenantMember{}).
|
||||
Where("organization_id = ? AND tenant_id = ?", orgID, tenantID).
|
||||
Update("role", role)
|
||||
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected == 0 {
|
||||
return ErrOrgMemberNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListTenantMembers lists all tenant memberships for an organization.
|
||||
func (r *organizationRepository) ListTenantMembers(ctx context.Context, orgID string) ([]*types.OrganizationTenantMember, error) {
|
||||
var members []*types.OrganizationTenantMember
|
||||
err := r.db.WithContext(ctx).
|
||||
Preload("RepresentativeUser").
|
||||
Where("organization_id = ?", orgID).
|
||||
Order("created_at ASC").
|
||||
Find(&members).Error
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return members, nil
|
||||
}
|
||||
|
||||
// GetTenantMember returns the (org, tenant) membership row, or
|
||||
// ErrOrgMemberNotFound when missing. This is the canonical permission
|
||||
// check primitive — callers compose this with the share permission
|
||||
// to compute effective access.
|
||||
func (r *organizationRepository) GetTenantMember(ctx context.Context, orgID string, tenantID uint64) (*types.OrganizationTenantMember, error) {
|
||||
var member types.OrganizationTenantMember
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("organization_id = ? AND tenant_id = ?", orgID, tenantID).
|
||||
First(&member).Error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrOrgMemberNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &member, nil
|
||||
}
|
||||
|
||||
// ListTenantMembersByTenantForOrgs returns one membership row per org where
|
||||
// the tenant participates (batch).
|
||||
func (r *organizationRepository) ListTenantMembersByTenantForOrgs(ctx context.Context, tenantID uint64, orgIDs []string) (map[string]*types.OrganizationTenantMember, error) {
|
||||
if len(orgIDs) == 0 {
|
||||
return make(map[string]*types.OrganizationTenantMember), nil
|
||||
}
|
||||
var members []*types.OrganizationTenantMember
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("tenant_id = ? AND organization_id IN ?", tenantID, orgIDs).
|
||||
Find(&members).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]*types.OrganizationTenantMember, len(members))
|
||||
for _, m := range members {
|
||||
if m != nil {
|
||||
out[m.OrganizationID] = m
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// CountTenantMembers counts the number of tenant members in an organization.
|
||||
func (r *organizationRepository) CountTenantMembers(ctx context.Context, orgID string) (int64, error) {
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).
|
||||
Model(&types.OrganizationTenantMember{}).
|
||||
Where("organization_id = ?", orgID).
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// UpdateInviteCode updates the invite code and optional expiry for an organization (expiresAt nil = never expire)
|
||||
func (r *organizationRepository) UpdateInviteCode(ctx context.Context, orgID string, inviteCode string, expiresAt *time.Time) error {
|
||||
updates := map[string]interface{}{"invite_code": inviteCode, "invite_code_expires_at": expiresAt}
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.Organization{}).
|
||||
Where("id = ?", orgID).
|
||||
Updates(updates).Error
|
||||
}
|
||||
|
||||
// ----------------
|
||||
// Join Requests
|
||||
// ----------------
|
||||
|
||||
var ErrJoinRequestNotFound = errors.New("join request not found")
|
||||
|
||||
// CreateJoinRequest creates a new join request
|
||||
func (r *organizationRepository) CreateJoinRequest(ctx context.Context, request *types.OrganizationJoinRequest) error {
|
||||
return r.db.WithContext(ctx).Create(request).Error
|
||||
}
|
||||
|
||||
// GetJoinRequestByID gets a join request by ID
|
||||
func (r *organizationRepository) GetJoinRequestByID(ctx context.Context, id string) (*types.OrganizationJoinRequest, error) {
|
||||
var request types.OrganizationJoinRequest
|
||||
err := r.db.WithContext(ctx).
|
||||
Preload("User").
|
||||
Where("id = ?", id).
|
||||
First(&request).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrJoinRequestNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &request, nil
|
||||
}
|
||||
|
||||
// GetPendingJoinRequestByTenant returns a pending request for the given
|
||||
// (org, tenant). After Plan 3 the dedup key is per-tenant, not per-user
|
||||
// — the user requesting and the user reviewing may be different humans
|
||||
// in the same tenant, but the resource being granted is tenant-level
|
||||
// access.
|
||||
func (r *organizationRepository) GetPendingJoinRequestByTenant(ctx context.Context, orgID string, tenantID uint64) (*types.OrganizationJoinRequest, error) {
|
||||
var request types.OrganizationJoinRequest
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("organization_id = ? AND tenant_id = ? AND status = ?", orgID, tenantID, types.JoinRequestStatusPending).
|
||||
First(&request).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrJoinRequestNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &request, nil
|
||||
}
|
||||
|
||||
// GetPendingRequestByTenantAndType narrows the (org, tenant) pending
|
||||
// dedup query to a specific request_type (join | upgrade).
|
||||
func (r *organizationRepository) GetPendingRequestByTenantAndType(ctx context.Context, orgID string, tenantID uint64, requestType types.JoinRequestType) (*types.OrganizationJoinRequest, error) {
|
||||
var request types.OrganizationJoinRequest
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("organization_id = ? AND tenant_id = ? AND status = ? AND request_type = ?", orgID, tenantID, types.JoinRequestStatusPending, requestType).
|
||||
First(&request).Error
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, ErrJoinRequestNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return &request, nil
|
||||
}
|
||||
|
||||
// ListJoinRequests lists join requests for an organization
|
||||
func (r *organizationRepository) ListJoinRequests(ctx context.Context, orgID string, status types.JoinRequestStatus) ([]*types.OrganizationJoinRequest, error) {
|
||||
var requests []*types.OrganizationJoinRequest
|
||||
query := r.db.WithContext(ctx).
|
||||
Preload("User").
|
||||
Where("organization_id = ?", orgID)
|
||||
|
||||
if status != "" {
|
||||
query = query.Where("status = ?", status)
|
||||
}
|
||||
|
||||
err := query.Order("created_at DESC").Find(&requests).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return requests, nil
|
||||
}
|
||||
|
||||
// CountJoinRequests counts join requests for an organization by status
|
||||
func (r *organizationRepository) CountJoinRequests(ctx context.Context, orgID string, status types.JoinRequestStatus) (int64, error) {
|
||||
var count int64
|
||||
query := r.db.WithContext(ctx).Model(&types.OrganizationJoinRequest{}).Where("organization_id = ?", orgID)
|
||||
if status != "" {
|
||||
query = query.Where("status = ?", status)
|
||||
}
|
||||
err := query.Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
// UpdateJoinRequestStatus updates the status of a join request
|
||||
func (r *organizationRepository) UpdateJoinRequestStatus(ctx context.Context, id string, status types.JoinRequestStatus, reviewedBy string, reviewMessage string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Model(&types.OrganizationJoinRequest{}).
|
||||
Where("id = ?", id).
|
||||
Updates(map[string]interface{}{
|
||||
"status": status,
|
||||
"reviewed_by": reviewedBy,
|
||||
"reviewed_at": gorm.Expr("NOW()"),
|
||||
"review_message": reviewMessage,
|
||||
}).Error
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
type resourceRepository struct{ db *gorm.DB }
|
||||
|
||||
// NewResourceRepository creates the persistence adapter for resource metadata.
|
||||
func NewResourceRepository(db *gorm.DB) interfaces.ResourceRepository {
|
||||
return &resourceRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *resourceRepository) Create(ctx context.Context, resource *types.StoredResource) error {
|
||||
return r.db.WithContext(ctx).Create(resource).Error
|
||||
}
|
||||
|
||||
func (r *resourceRepository) GetByID(ctx context.Context, id string) (*types.StoredResource, error) {
|
||||
var resource types.StoredResource
|
||||
err := r.db.WithContext(ctx).Where("id = ? AND state = ?", id, types.ResourceStateActive).First(&resource).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return &resource, err
|
||||
}
|
||||
|
||||
func (r *resourceRepository) GetByHandle(ctx context.Context, handle string) (*types.StoredResource, error) {
|
||||
var resource types.StoredResource
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("handle = ? AND state = ?", handle, types.ResourceStateActive).
|
||||
First(&resource).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return &resource, err
|
||||
}
|
||||
|
||||
func (r *resourceRepository) GetByTenantLocation(
|
||||
ctx context.Context,
|
||||
tenantID uint64,
|
||||
locationHash string,
|
||||
) (*types.StoredResource, error) {
|
||||
var resource types.StoredResource
|
||||
err := r.db.WithContext(ctx).
|
||||
Where(
|
||||
"tenant_id = ? AND location_hash = ? AND state = ?",
|
||||
tenantID,
|
||||
locationHash,
|
||||
types.ResourceStateActive,
|
||||
).
|
||||
First(&resource).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return &resource, err
|
||||
}
|
||||
|
||||
func (r *resourceRepository) MarkDeleted(ctx context.Context, id string) error {
|
||||
return r.db.WithContext(ctx).Model(&types.StoredResource{}).Where("id = ?", id).
|
||||
Updates(map[string]interface{}{"state": types.ResourceStateDeleted, "deleted_at": time.Now()}).Error
|
||||
}
|
||||
|
||||
func (r *resourceRepository) CreateBinding(ctx context.Context, binding *types.ResourceBinding) error {
|
||||
return r.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(binding).Error
|
||||
}
|
||||
|
||||
// DeleteBinding removes one owner's claim on a resource. Deleting a claim that
|
||||
// was never recorded is not an error: callers release optimistically, from a
|
||||
// content scan that cannot know which references were bound.
|
||||
func (r *resourceRepository) DeleteBinding(ctx context.Context, resourceID, ownerType, ownerID string) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("resource_id = ? AND owner_type = ? AND owner_id = ?", resourceID, ownerType, ownerID).
|
||||
Delete(&types.ResourceBinding{}).Error
|
||||
}
|
||||
|
||||
func (r *resourceRepository) CountBindings(ctx context.Context, resourceID string) (int64, error) {
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).Model(&types.ResourceBinding{}).
|
||||
Where("resource_id = ?", resourceID).
|
||||
Count(&count).Error
|
||||
return count, err
|
||||
}
|
||||
|
||||
func (r *resourceRepository) CreateGrant(ctx context.Context, grant *types.ResourceAccessGrant) error {
|
||||
return r.db.WithContext(ctx).Create(grant).Error
|
||||
}
|
||||
|
||||
func (r *resourceRepository) GetValidGrant(
|
||||
ctx context.Context,
|
||||
tokenHash string,
|
||||
now time.Time,
|
||||
) (*types.ResourceAccessGrant, error) {
|
||||
var grant types.ResourceAccessGrant
|
||||
err := r.db.WithContext(ctx).
|
||||
Where("token_hash = ? AND revoked_at IS NULL AND expires_at > ?", tokenHash, now).
|
||||
First(&grant).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, nil
|
||||
}
|
||||
return &grant, err
|
||||
}
|
||||
|
||||
// DeleteExpiredGrants drops grants that are past their expiry. A revoked grant
|
||||
// is kept until then on purpose: it is the tombstone that stops a token derived
|
||||
// for the same resource and window from re-creating the row and reviving access.
|
||||
func (r *resourceRepository) DeleteExpiredGrants(ctx context.Context, before time.Time) error {
|
||||
return r.db.WithContext(ctx).
|
||||
Where("expires_at <= ?", before).
|
||||
Delete(&types.ResourceAccessGrant{}).Error
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var errFileAccessDenied = errors.New("file is not attached to an accessible resource")
|
||||
|
||||
// A storage locator is never an authorization capability. Resolve it back to
|
||||
// live domain owners, then apply the same personal/knowledge scope as their APIs.
|
||||
func (r *resourceRepository) AuthorizeFileRead(ctx context.Context, value string, scope interfaces.ResourceFileScope) error {
|
||||
if scope.TenantID == 0 {
|
||||
return errFileAccessDenied
|
||||
}
|
||||
var resource *types.StoredResource
|
||||
var err error
|
||||
if handle, ok := types.ParseResourcePath(value); ok {
|
||||
resource, err = r.GetByHandle(ctx, handle)
|
||||
} else {
|
||||
hash := sha256.Sum256([]byte(value))
|
||||
resource, err = r.GetByTenantLocation(ctx, scope.TenantID, hex.EncodeToString(hash[:]))
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var bindings []types.ResourceBinding
|
||||
if resource != nil {
|
||||
if resource.TenantID != scope.TenantID {
|
||||
return errFileAccessDenied
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Where("resource_id = ? AND tenant_id = ?", resource.ID, scope.TenantID).Find(&bindings).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// Older registrations may lack a source_file relation. An attachment
|
||||
// alias for the same original must not bypass the download restriction.
|
||||
var originals []types.Knowledge
|
||||
if err := r.db.WithContext(ctx).Select("id").Where("tenant_id = ? AND file_path IN ?", scope.TenantID,
|
||||
[]string{value, resource.PhysicalPath, types.BuildResourcePath(resource.Handle)}).Find(&originals).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, doc := range originals {
|
||||
bindings = append(bindings, types.ResourceBinding{OwnerType: types.ResourceOwnerKnowledge, OwnerID: doc.ID, Relation: types.ResourceRelationSourceFile})
|
||||
}
|
||||
} else if _, isHandle := types.ParseResourcePath(value); !isHandle {
|
||||
// Legacy source files predate resource bindings but still have an exact
|
||||
// knowledge owner. Do not accept an arbitrary same-tenant path.
|
||||
var documents []types.Knowledge
|
||||
if err := r.db.WithContext(ctx).Select("id").Where("tenant_id = ? AND file_path = ?", scope.TenantID, value).Find(&documents).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, doc := range documents {
|
||||
bindings = append(bindings, types.ResourceBinding{OwnerType: types.ResourceOwnerKnowledge, OwnerID: doc.ID, Relation: types.ResourceRelationSourceFile})
|
||||
}
|
||||
}
|
||||
var sourceBindings []types.ResourceBinding
|
||||
for _, binding := range bindings {
|
||||
if binding.OwnerType == types.ResourceOwnerKnowledge && binding.Relation == types.ResourceRelationSourceFile {
|
||||
sourceBindings = append(sourceBindings, binding)
|
||||
}
|
||||
}
|
||||
if len(sourceBindings) > 0 {
|
||||
bindings = sourceBindings
|
||||
}
|
||||
for _, binding := range bindings {
|
||||
allowed, err := r.readableFileBinding(ctx, scope, binding)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if allowed {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return errFileAccessDenied
|
||||
}
|
||||
|
||||
func (r *resourceRepository) readableFileBinding(ctx context.Context, scope interfaces.ResourceFileScope, binding types.ResourceBinding) (bool, error) {
|
||||
q := r.db.WithContext(ctx)
|
||||
callerTenant, _ := types.TenantIDFromContext(ctx)
|
||||
userID, _ := types.UserIDFromContext(ctx)
|
||||
switch binding.OwnerType {
|
||||
case types.ResourceOwnerKnowledge:
|
||||
if binding.Relation == types.ResourceRelationSourceFile {
|
||||
if callerTenant != scope.TenantID || userID == "" || scope.EmbedChannelID != "" || !accesspolicy.Allowed(ctx, "knowledge.read") {
|
||||
return false, nil
|
||||
}
|
||||
var source types.Knowledge
|
||||
if err := q.Select("id,knowledge_base_id").Where("id = ? AND tenant_id = ?", binding.OwnerID, scope.TenantID).Take(&source).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
var base types.KnowledgeBase
|
||||
if err := q.Select("id,creator_id").Where("id = ? AND tenant_id = ?", source.KnowledgeBaseID, scope.TenantID).Take(&base).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
if base.CreatorID != userID && !types.IsSystemAdminFromContext(ctx) && !types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin) {
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
if scope.MessageID != "" {
|
||||
var message types.Message
|
||||
if err := q.Select("id,session_id,knowledge_references").Where("id = ?", scope.MessageID).Take(&message).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
for _, reference := range message.KnowledgeReferences {
|
||||
if reference != nil && reference.KnowledgeID == binding.OwnerID {
|
||||
var count int64
|
||||
if err := q.Model(&types.Knowledge{}).Where("id = ? AND tenant_id = ?", binding.OwnerID, scope.TenantID).Count(&count).Error; err != nil {
|
||||
return false, err
|
||||
}
|
||||
if count == 0 {
|
||||
return false, nil
|
||||
}
|
||||
return r.readableFileSession(ctx, scope, message.SessionID)
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
var knowledge types.Knowledge
|
||||
if err := q.Select("id,knowledge_base_id").Where("id = ? AND tenant_id = ?", binding.OwnerID, scope.TenantID).Take(&knowledge).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
if scope.KnowledgeBaseID != "" {
|
||||
// KBAccessRead has already authorized this exact knowledge base.
|
||||
return knowledge.KnowledgeBaseID == scope.KnowledgeBaseID, nil
|
||||
}
|
||||
var base types.KnowledgeBase
|
||||
if err := q.Select("id,tenant_id,creator_id,is_temporary,version_state").Where("id = ? AND tenant_id = ?", knowledge.KnowledgeBaseID, scope.TenantID).Take(&base).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
if base.IsTemporary || base.VersionState != "" {
|
||||
return false, nil
|
||||
}
|
||||
if scope.EmbedChannelID != "" {
|
||||
var channel types.EmbedChannel
|
||||
if err := q.Where("id = ? AND tenant_id = ? AND enabled = ?", scope.EmbedChannelID, scope.TenantID, true).Take(&channel).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
var agent types.CustomAgent
|
||||
if err := q.Where("id = ? AND tenant_id = ? AND enabled = ?", channel.AgentID, scope.TenantID, true).Take(&agent).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
switch agent.Config.KBSelectionMode {
|
||||
case "none":
|
||||
return false, nil
|
||||
case "selected":
|
||||
for _, id := range agent.Config.KnowledgeBases {
|
||||
if id == base.ID {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
return false, nil
|
||||
default:
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
if callerTenant != scope.TenantID || userID == "" || !accesspolicy.Allowed(ctx, "knowledge.read") {
|
||||
return false, nil
|
||||
}
|
||||
if types.IsSystemAdminFromContext(ctx) || types.TenantRoleFromContext(ctx).HasPermission(types.TenantRoleAdmin) {
|
||||
return true, nil
|
||||
}
|
||||
return accesspolicy.Readable(ctx, "knowledge", base.ID, base.CreatorID == userID, accesspolicy.FromContext(ctx) == nil), nil
|
||||
case types.ResourceOwnerMessage, types.ResourceOwnerTemporaryDocument:
|
||||
userID = types.SessionOwnerIDFromContext(ctx)
|
||||
if scope.KnowledgeBaseID != "" || userID == "" {
|
||||
return false, nil
|
||||
}
|
||||
var sessionID string
|
||||
if binding.OwnerType == types.ResourceOwnerMessage {
|
||||
if scope.MessageID != "" && binding.OwnerID != scope.MessageID {
|
||||
return false, nil
|
||||
}
|
||||
var message types.Message
|
||||
if err := q.Select("id,session_id").Where("id = ?", binding.OwnerID).Take(&message).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
sessionID = message.SessionID
|
||||
} else {
|
||||
if scope.MessageID != "" {
|
||||
return false, nil
|
||||
}
|
||||
var document types.TemporaryDocument
|
||||
if err := q.Select("session_id").Where("id = ? AND tenant_id = ? AND expires_at > ?", binding.OwnerID, scope.TenantID, time.Now()).Take(&document).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
sessionID = document.SessionID
|
||||
}
|
||||
return r.readableFileSession(ctx, scope, sessionID)
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
func (r *resourceRepository) readableFileSession(ctx context.Context, scope interfaces.ResourceFileScope, sessionID string) (bool, error) {
|
||||
tenantID, _ := types.TenantIDFromContext(ctx)
|
||||
userID := types.SessionOwnerIDFromContext(ctx)
|
||||
if userID == "" || scope.SessionID != "" && scope.SessionID != sessionID {
|
||||
return false, nil
|
||||
}
|
||||
var session types.Session
|
||||
if err := r.db.WithContext(ctx).Select("id,tenant_id,user_id").Where("id = ? AND tenant_id = ? AND user_id = ?", sessionID, tenantID, userID).Take(&session).Error; err != nil {
|
||||
return false, ignoreMissingFileOwner(err)
|
||||
}
|
||||
if scope.EmbedChannelID != "" {
|
||||
return scope.SessionID != "" && strings.HasPrefix(userID, types.PrincipalEmbedSession+":"), nil
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func ignoreMissingFileOwner(err error) error {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func fileAccessFixture(t *testing.T) (*resourceRepository, *gorm.DB, string) {
|
||||
t.Helper()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
require.NoError(t, db.AutoMigrate(&types.StoredResource{}, &types.ResourceBinding{}))
|
||||
for _, schema := range []string{
|
||||
"CREATE TABLE knowledges (id TEXT PRIMARY KEY, tenant_id INTEGER, knowledge_base_id TEXT, file_path TEXT, deleted_at DATETIME)",
|
||||
"CREATE TABLE knowledge_bases (id TEXT PRIMARY KEY, tenant_id INTEGER, creator_id TEXT, is_temporary BOOLEAN DEFAULT FALSE, version_state TEXT DEFAULT '', deleted_at DATETIME)",
|
||||
"CREATE TABLE sessions (id TEXT PRIMARY KEY, tenant_id INTEGER, user_id TEXT, description TEXT, deleted_at DATETIME)",
|
||||
"CREATE TABLE messages (id TEXT PRIMARY KEY, session_id TEXT, knowledge_references TEXT, deleted_at DATETIME)",
|
||||
"CREATE TABLE temporary_documents (id TEXT PRIMARY KEY, tenant_id INTEGER, session_id TEXT, expires_at DATETIME, deleted_at DATETIME)",
|
||||
"CREATE TABLE custom_agents (id TEXT PRIMARY KEY, tenant_id INTEGER, enabled BOOLEAN, config TEXT, deleted_at DATETIME)",
|
||||
"CREATE TABLE embed_channels (id TEXT PRIMARY KEY, tenant_id INTEGER, agent_id TEXT, enabled BOOLEAN, deleted_at DATETIME)",
|
||||
} {
|
||||
require.NoError(t, db.Exec(schema).Error)
|
||||
}
|
||||
resource := &types.StoredResource{ID: "file-1", Handle: "AbCdEfGhIjKlMnOpQrStUv", TenantID: 7, PhysicalPath: "local://7/exports/private.png", Provider: "local", LocationHash: "unused"}
|
||||
require.NoError(t, db.Create(resource).Error)
|
||||
return &resourceRepository{db: db}, db, types.BuildResourcePath(resource.Handle)
|
||||
}
|
||||
|
||||
func fileUserContext(user string) context.Context {
|
||||
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
|
||||
ctx = context.WithValue(ctx, types.UserIDContextKey, user)
|
||||
return accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": false}})
|
||||
}
|
||||
|
||||
func TestResourceFileAccessKnowledgeBinding(t *testing.T) {
|
||||
r, db, ref := fileAccessFixture(t)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice'),('kb-b',7,'bob')").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id,file_path) VALUES ('doc-a',7,'kb-a','local://7/legacy.doc')").Error)
|
||||
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
|
||||
scope := interfaces.ResourceFileScope{TenantID: 7}
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), "local://7/legacy.doc", scope))
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), "local://7/legacy.doc", scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), "local://7/unregistered.png", scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 8}))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-b"}))
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}))
|
||||
shared := accesspolicy.WithSnapshot(fileUserContext("bob"), &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": false}, Grants: map[string]bool{"knowledge:kb-a": true}})
|
||||
require.NoError(t, r.AuthorizeFileRead(shared, ref, scope))
|
||||
denied := accesspolicy.WithSnapshot(fileUserContext("alice"), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{}})
|
||||
require.Error(t, r.AuthorizeFileRead(denied, ref, scope))
|
||||
require.NoError(t, db.Exec("UPDATE knowledges SET deleted_at=CURRENT_TIMESTAMP WHERE id='doc-a'").Error)
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
}
|
||||
|
||||
func TestResourceFileAccessSharedOriginalCannotBeDownloaded(t *testing.T) {
|
||||
r, db, ref := fileAccessFixture(t)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice')").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id,file_path) VALUES ('doc-a',7,'kb-a','local://7/legacy.doc')").Error)
|
||||
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a", Relation: types.ResourceRelationSourceFile}))
|
||||
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a", Relation: types.ResourceRelationAttachment}))
|
||||
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-b',7,'bob')").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id,knowledge_references) VALUES ('msg-b','s-b',?)", []byte(`[{"knowledge_id":"doc-a"}]`)).Error)
|
||||
shared := accesspolicy.WithSnapshot(fileUserContext("bob"), &accesspolicy.Snapshot{Publications: map[string]bool{"knowledge:kb-a": true}, Grants: map[string]bool{"knowledge:kb-a": true}})
|
||||
for _, scope := range []interfaces.ResourceFileScope{
|
||||
{TenantID: 7},
|
||||
{TenantID: 7, KnowledgeBaseID: "kb-a"},
|
||||
{TenantID: 7, MessageID: "msg-b", SessionID: "s-b"},
|
||||
{TenantID: 7, EmbedChannelID: "channel-a"},
|
||||
} {
|
||||
require.ErrorIs(t, r.AuthorizeFileRead(shared, ref, scope), errFileAccessDenied)
|
||||
require.ErrorIs(t, r.AuthorizeFileRead(shared, "local://7/legacy.doc", scope), errFileAccessDenied)
|
||||
}
|
||||
scope := interfaces.ResourceFileScope{TenantID: 7}
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
admin := context.WithValue(fileUserContext("admin"), types.TenantRoleContextKey, types.TenantRoleAdmin)
|
||||
require.NoError(t, r.AuthorizeFileRead(admin, ref, scope))
|
||||
systemAdmin := context.WithValue(fileUserContext("system"), types.SystemAdminContextKey, true)
|
||||
require.NoError(t, r.AuthorizeFileRead(systemAdmin, ref, scope))
|
||||
crossTenant := context.WithValue(systemAdmin, types.TenantIDContextKey, uint64(8))
|
||||
require.ErrorIs(t, r.AuthorizeFileRead(crossTenant, ref, scope), errFileAccessDenied)
|
||||
denied := accesspolicy.WithSnapshot(fileUserContext("alice"), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{}})
|
||||
require.ErrorIs(t, r.AuthorizeFileRead(denied, ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}), errFileAccessDenied)
|
||||
|
||||
// A legacy registration without a source relation is still an original.
|
||||
require.NoError(t, db.Exec("DELETE FROM resource_bindings WHERE relation=?", types.ResourceRelationSourceFile).Error)
|
||||
for _, location := range []string{ref, "local://7/exports/private.png"} {
|
||||
require.NoError(t, db.Exec("UPDATE knowledges SET file_path=? WHERE id='doc-a'", location).Error)
|
||||
require.ErrorIs(t, r.AuthorizeFileRead(shared, ref, scope), errFileAccessDenied)
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
}
|
||||
}
|
||||
|
||||
func TestResourceFileAccessPrivateMessages(t *testing.T) {
|
||||
r, db, ref := fileAccessFixture(t)
|
||||
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,'alice'),('s-b',7,'bob'),('s-legacy',7,'')").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id) VALUES ('msg-a','s-a'),('msg-b','s-b'),('msg-legacy','s-legacy')").Error)
|
||||
binding := &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerMessage, OwnerID: "msg-a"}
|
||||
require.NoError(t, r.CreateBinding(context.Background(), binding))
|
||||
scope := interfaces.ResourceFileScope{TenantID: 7}
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-b", SessionID: "s-b"}))
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-a", SessionID: "s-a"}))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, interfaces.ResourceFileScope{TenantID: 7, KnowledgeBaseID: "kb-a"}))
|
||||
require.NoError(t, db.Model(&types.ResourceBinding{}).Where("id = ?", binding.ID).Update("owner_id", "msg-legacy").Error)
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||||
}
|
||||
|
||||
func TestResourceFileAccessEmbedKnowledgeScope(t *testing.T) {
|
||||
r, db, ref := fileAccessFixture(t)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledge_bases (id,tenant_id,creator_id) VALUES ('kb-a',7,'alice')").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id) VALUES ('doc-a',7,'kb-a')").Error)
|
||||
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
|
||||
require.NoError(t, db.Exec("INSERT INTO custom_agents (id,tenant_id,enabled,config) VALUES ('agent-a',7,TRUE,?)", `{"kb_selection_mode":"selected","knowledge_bases":["kb-a"]}`).Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO embed_channels (id,tenant_id,agent_id,enabled) VALUES ('channel-a',7,'agent-a',TRUE)").Error)
|
||||
scope := interfaces.ResourceFileScope{TenantID: 7, EmbedChannelID: "channel-a"}
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("visitor"), ref, scope))
|
||||
for _, config := range []string{`{"kb_selection_mode":"none"}`, `{"kb_selection_mode":"selected","knowledge_bases":["kb-other"]}`, `{"kb_selection_mode":"selected","knowledge_bases":[]}`} {
|
||||
require.NoError(t, db.Exec("UPDATE custom_agents SET config=?", config).Error)
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("visitor"), ref, scope))
|
||||
}
|
||||
}
|
||||
|
||||
func TestResourceFileAccessMessageKnowledgeReferences(t *testing.T) {
|
||||
r, db, ref := fileAccessFixture(t)
|
||||
require.NoError(t, db.Exec("INSERT INTO knowledges (id,tenant_id,knowledge_base_id) VALUES ('doc-a',7,'kb-a')").Error)
|
||||
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerKnowledge, OwnerID: "doc-a"}))
|
||||
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,'alice')").Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id,knowledge_references) VALUES ('msg-a','s-a',?)", []byte(`[{"knowledge_id":"doc-a"}]`)).Error)
|
||||
scope := interfaces.ResourceFileScope{TenantID: 7, MessageID: "msg-a", SessionID: "s-a"}
|
||||
require.NoError(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("bob"), ref, scope))
|
||||
require.NoError(t, db.Exec("UPDATE messages SET knowledge_references=?", []byte(`[{"knowledge_id":"doc-unrelated"}]`)).Error)
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("alice"), ref, scope))
|
||||
}
|
||||
|
||||
func TestResourceFileAccessEmbedPrivateSession(t *testing.T) {
|
||||
r, db, ref := fileAccessFixture(t)
|
||||
principal := types.EmbedSessionPrincipal(7, "channel-a", "s-a")
|
||||
require.NoError(t, db.Exec("INSERT INTO sessions (id,tenant_id,user_id) VALUES ('s-a',7,?)", principal.StorageID()).Error)
|
||||
require.NoError(t, db.Exec("INSERT INTO messages (id,session_id) VALUES ('msg-a','s-a')").Error)
|
||||
require.NoError(t, r.CreateBinding(context.Background(), &types.ResourceBinding{ResourceID: "file-1", TenantID: 7, OwnerType: types.ResourceOwnerMessage, OwnerID: "msg-a"}))
|
||||
ctx := types.WithPrincipal(fileUserContext("embed"), principal)
|
||||
scope := interfaces.ResourceFileScope{TenantID: 7, EmbedChannelID: "channel-a", SessionID: "s-a"}
|
||||
require.NoError(t, r.AuthorizeFileRead(ctx, ref, scope))
|
||||
require.Error(t, r.AuthorizeFileRead(fileUserContext("embed"), ref, scope))
|
||||
scope.SessionID = "s-other"
|
||||
require.Error(t, r.AuthorizeFileRead(ctx, ref, scope))
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
const envDorisCompatMode = "DORIS_COMPAT_MODE"
|
||||
|
||||
type dorisCompatMode string
|
||||
|
||||
const (
|
||||
dorisCompatModeAuto dorisCompatMode = "auto"
|
||||
dorisCompatModeLegacy dorisCompatMode = "legacy"
|
||||
dorisCompatModeInnerProductDuplicate dorisCompatMode = "inner_product_duplicate"
|
||||
)
|
||||
|
||||
type dorisCompatProbe struct {
|
||||
innerProductApproximate bool
|
||||
cosineDistanceApproximate bool
|
||||
}
|
||||
|
||||
func resolveConfiguredDorisCompatMode() (mode dorisCompatMode, invalid string) {
|
||||
raw := strings.TrimSpace(os.Getenv(envDorisCompatMode))
|
||||
if raw == "" {
|
||||
return dorisCompatModeAuto, ""
|
||||
}
|
||||
|
||||
switch strings.ToLower(raw) {
|
||||
case string(dorisCompatModeAuto):
|
||||
return dorisCompatModeAuto, ""
|
||||
case string(dorisCompatModeLegacy):
|
||||
return dorisCompatModeLegacy, ""
|
||||
case string(dorisCompatModeInnerProductDuplicate), "inner-product-duplicate", "inner_product", "inner-product":
|
||||
return dorisCompatModeInnerProductDuplicate, ""
|
||||
default:
|
||||
return dorisCompatModeAuto, raw
|
||||
}
|
||||
}
|
||||
|
||||
func (r *dorisRepository) resolveCompatMode(ctx context.Context) (dorisCompatMode, error) {
|
||||
r.compatResolveOnce.Do(func() {
|
||||
log := logger.GetLogger(ctx)
|
||||
requested := r.compatModeRequested
|
||||
if requested == "" {
|
||||
r.compatModeResolved = dorisCompatModeInnerProductDuplicate
|
||||
return
|
||||
}
|
||||
|
||||
existingMode, exampleTable, foundExisting, err := r.detectExistingCompatMode(ctx)
|
||||
if err != nil {
|
||||
r.compatResolveErr = err
|
||||
return
|
||||
}
|
||||
|
||||
if foundExisting {
|
||||
if requested != dorisCompatModeAuto && requested != existingMode {
|
||||
r.compatResolveErr = fmt.Errorf(
|
||||
"Doris compat mode %q does not match existing embedding tables (detected %q from %s). %s is not interchangeable after %s_* tables are created. Recreate the existing %s_* tables before switching modes, or set %s=%s",
|
||||
requested,
|
||||
existingMode,
|
||||
exampleTable,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
r.tableBaseName,
|
||||
envDorisCompatMode,
|
||||
existingMode,
|
||||
)
|
||||
log.Errorf("[Doris] %v", r.compatResolveErr)
|
||||
return
|
||||
}
|
||||
|
||||
r.compatModeResolved = existingMode
|
||||
log.Warnf(
|
||||
"[Doris] Using compat mode %s from existing embedding tables (%s). %s is not interchangeable after %s_* tables are created; recreate those tables before switching modes.",
|
||||
existingMode,
|
||||
exampleTable,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
if requested == dorisCompatModeAuto {
|
||||
probe, probeErr := r.probeCompatMode(ctx)
|
||||
if probeErr != nil {
|
||||
r.compatResolveErr = probeErr
|
||||
log.Errorf("[Doris] %v", probeErr)
|
||||
return
|
||||
}
|
||||
log.Infof(
|
||||
"[Doris] Compat probe result: inner_product_approximate=%t, cosine_distance_approximate=%t",
|
||||
probe.innerProductApproximate,
|
||||
probe.cosineDistanceApproximate,
|
||||
)
|
||||
|
||||
switch {
|
||||
case probe.innerProductApproximate:
|
||||
r.compatModeResolved = dorisCompatModeInnerProductDuplicate
|
||||
case probe.cosineDistanceApproximate:
|
||||
r.compatModeResolved = dorisCompatModeLegacy
|
||||
default:
|
||||
r.compatResolveErr = fmt.Errorf(
|
||||
"Doris compatibility auto-detection could not find a supported vector function. Set %s=%s or %s=%s explicitly after verifying your Doris build. %s is not interchangeable after %s_* tables are created",
|
||||
envDorisCompatMode,
|
||||
dorisCompatModeInnerProductDuplicate,
|
||||
envDorisCompatMode,
|
||||
dorisCompatModeLegacy,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
)
|
||||
log.Errorf("[Doris] %v", r.compatResolveErr)
|
||||
return
|
||||
}
|
||||
|
||||
log.Warnf(
|
||||
"[Doris] Auto-selected compat mode %s for new embedding tables. %s is not interchangeable after %s_* tables are created; recreate those tables before switching modes.",
|
||||
r.compatModeResolved,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
r.compatModeResolved = requested
|
||||
log.Warnf(
|
||||
"[Doris] Using configured compat mode %s for new embedding tables. %s is not interchangeable after %s_* tables are created; recreate those tables before switching modes.",
|
||||
r.compatModeResolved,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
)
|
||||
})
|
||||
|
||||
return r.compatModeResolved, r.compatResolveErr
|
||||
}
|
||||
|
||||
func (r *dorisRepository) detectExistingCompatMode(ctx context.Context) (dorisCompatMode, string, bool, error) {
|
||||
tables, err := r.listEmbeddingTables(ctx)
|
||||
if err != nil {
|
||||
return "", "", false, fmt.Errorf("list Doris embedding tables: %w", err)
|
||||
}
|
||||
if len(tables) == 0 {
|
||||
return "", "", false, nil
|
||||
}
|
||||
sort.Strings(tables)
|
||||
|
||||
var detected dorisCompatMode
|
||||
for _, table := range tables {
|
||||
ddl, err := r.showCreateTable(ctx, table)
|
||||
if err != nil {
|
||||
return "", "", false, fmt.Errorf("show create table %s: %w", table, err)
|
||||
}
|
||||
mode, err := compatModeFromDDL(ddl)
|
||||
if err != nil {
|
||||
return "", "", false, fmt.Errorf("detect compat mode from %s: %w", table, err)
|
||||
}
|
||||
if detected == "" {
|
||||
detected = mode
|
||||
continue
|
||||
}
|
||||
if detected != mode {
|
||||
return "", "", false, fmt.Errorf(
|
||||
"existing Doris embedding tables use mixed compat modes (%s and %s). %s is not interchangeable after %s_* tables are created; recreate the existing %s_* tables with a single mode",
|
||||
detected,
|
||||
mode,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
r.tableBaseName,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
return detected, tables[0], true, nil
|
||||
}
|
||||
|
||||
func (r *dorisRepository) showCreateTable(ctx context.Context, table string) (string, error) {
|
||||
row := r.db.QueryRowContext(ctx, fmt.Sprintf("SHOW CREATE TABLE `%s`", table))
|
||||
var tableName, ddl string
|
||||
if err := row.Scan(&tableName, &ddl); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return ddl, nil
|
||||
}
|
||||
|
||||
func compatModeFromDDL(ddl string) (dorisCompatMode, error) {
|
||||
normalized := strings.ToLower(ddl)
|
||||
switch {
|
||||
case strings.Contains(normalized, "duplicate key("):
|
||||
return dorisCompatModeInnerProductDuplicate, nil
|
||||
case strings.Contains(normalized, "unique key("):
|
||||
return dorisCompatModeLegacy, nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported table definition; expected UNIQUE KEY or DUPLICATE KEY")
|
||||
}
|
||||
}
|
||||
|
||||
func (r *dorisRepository) probeCompatMode(ctx context.Context) (dorisCompatProbe, error) {
|
||||
probe := dorisCompatProbe{}
|
||||
probe.innerProductApproximate = r.vectorFunctionSupported(ctx, "inner_product_approximate([1.0],[1.0])")
|
||||
probe.cosineDistanceApproximate = r.vectorFunctionSupported(ctx, "cosine_distance_approximate([1.0],[1.0])")
|
||||
return probe, nil
|
||||
}
|
||||
|
||||
func (r *dorisRepository) vectorFunctionSupported(ctx context.Context, expr string) bool {
|
||||
var value sql.NullFloat64
|
||||
if err := r.db.QueryRowContext(ctx, "SELECT "+expr).Scan(&value); err != nil {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (m dorisCompatMode) normalizeEmbeddings() bool {
|
||||
return m != dorisCompatModeLegacy
|
||||
}
|
||||
|
||||
func (m dorisCompatMode) usesReplaceWrite() bool {
|
||||
return m != dorisCompatModeLegacy
|
||||
}
|
||||
|
||||
func (m dorisCompatMode) usesRewriteChunkUpdates() bool {
|
||||
return m != dorisCompatModeLegacy
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// 字段名常量。Doris 是 SQL 库,字段名要在 SELECT/WHERE/INSERT 多处复用,
|
||||
// 用常量统一防止笔误。
|
||||
const (
|
||||
fieldID = "id"
|
||||
fieldContent = "content"
|
||||
fieldSourceID = "source_id"
|
||||
fieldSourceType = "source_type"
|
||||
fieldChunkID = "chunk_id"
|
||||
fieldKnowledgeID = "knowledge_id"
|
||||
fieldKnowledgeBaseID = "knowledge_base_id"
|
||||
fieldTagID = "tag_id"
|
||||
fieldIsEnabled = "is_enabled"
|
||||
fieldEmbedding = "embedding"
|
||||
)
|
||||
|
||||
// columns 是 INSERT / SELECT 时使用的标准列序。
|
||||
var columns = []string{
|
||||
fieldID, fieldContent, fieldSourceID, fieldSourceType,
|
||||
fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID,
|
||||
fieldIsEnabled, fieldEmbedding,
|
||||
}
|
||||
|
||||
// columnsForRetrieve 是 Retrieve 时 SELECT 的列序,
|
||||
// 不包含 embedding(向量本身查询结果中无需返回,省带宽)。
|
||||
var columnsForRetrieve = []string{
|
||||
fieldID, fieldContent, fieldSourceID, fieldSourceType,
|
||||
fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID,
|
||||
fieldIsEnabled,
|
||||
}
|
||||
|
||||
// columnsForCopy 是 CopyIndices 中分页 SELECT 时使用的列序,
|
||||
// 比 columnsForRetrieve 多 embedding,因为复制目的是搬运向量本身。
|
||||
var columnsForCopy = []string{
|
||||
fieldID, fieldContent, fieldSourceID, fieldSourceType,
|
||||
fieldChunkID, fieldKnowledgeID, fieldKnowledgeBaseID, fieldTagID,
|
||||
fieldIsEnabled, fieldEmbedding,
|
||||
}
|
||||
|
||||
// whereCond 表示一个 WHERE 子条件:clause 是参数化 SQL 片段(带 ? 占位),
|
||||
// args 是对应顺序的参数值。所有用户输入字段(IDs)必须通过 args 传入,
|
||||
// 严禁拼到 clause 字符串里。
|
||||
type whereCond struct {
|
||||
clause string
|
||||
args []any
|
||||
}
|
||||
|
||||
// whereBuilder 用于把 RetrieveParams 中的过滤条件翻译成 SQL WHERE 子句。
|
||||
//
|
||||
// 每个 add* 方法对应一种 IN / NOT IN / = 算子;最终 build() 用 AND 拼接。
|
||||
type whereBuilder struct {
|
||||
conds []whereCond
|
||||
}
|
||||
|
||||
// addEqual 追加一个 field = ? 条件。
|
||||
func (w *whereBuilder) addEqual(field string, value any) {
|
||||
w.conds = append(w.conds, whereCond{
|
||||
clause: field + " = ?",
|
||||
args: []any{value},
|
||||
})
|
||||
}
|
||||
|
||||
// addIn 追加一个 field IN (?, ?, ...) 条件。values 为空时不追加任何东西。
|
||||
func (w *whereBuilder) addIn(field string, values []string) {
|
||||
if len(values) == 0 {
|
||||
return
|
||||
}
|
||||
placeholders := make([]string, len(values))
|
||||
args := make([]any, len(values))
|
||||
for i, v := range values {
|
||||
placeholders[i] = "?"
|
||||
args[i] = v
|
||||
}
|
||||
w.conds = append(w.conds, whereCond{
|
||||
clause: field + " IN (" + strings.Join(placeholders, ", ") + ")",
|
||||
args: args,
|
||||
})
|
||||
}
|
||||
|
||||
// addNotIn 追加一个 field NOT IN (?, ?, ...) 条件。
|
||||
func (w *whereBuilder) addNotIn(field string, values []string) {
|
||||
if len(values) == 0 {
|
||||
return
|
||||
}
|
||||
placeholders := make([]string, len(values))
|
||||
args := make([]any, len(values))
|
||||
for i, v := range values {
|
||||
placeholders[i] = "?"
|
||||
args[i] = v
|
||||
}
|
||||
w.conds = append(w.conds, whereCond{
|
||||
clause: field + " NOT IN (" + strings.Join(placeholders, ", ") + ")",
|
||||
args: args,
|
||||
})
|
||||
}
|
||||
|
||||
// build 返回 WHERE 子句(不含 "WHERE " 前缀)和参数数组。
|
||||
// 没有任何条件时返回 ("1 = 1", nil),方便调用方无脑拼接。
|
||||
func (w *whereBuilder) build() (string, []any) {
|
||||
if len(w.conds) == 0 {
|
||||
return "1 = 1", nil
|
||||
}
|
||||
parts := make([]string, len(w.conds))
|
||||
var args []any
|
||||
for i, c := range w.conds {
|
||||
parts[i] = c.clause
|
||||
args = append(args, c.args...)
|
||||
}
|
||||
return strings.Join(parts, " AND "), args
|
||||
}
|
||||
|
||||
// buildBaseFilter 将 RetrieveParams 中的过滤条件翻译为 whereBuilder。
|
||||
// 默认追加 is_enabled = TRUE,与 Qdrant/Milvus/Weaviate 保持一致:
|
||||
// 关闭的 chunk 不参与检索。
|
||||
func buildBaseFilter(params types.RetrieveParams) *whereBuilder {
|
||||
w := &whereBuilder{}
|
||||
w.addEqual(fieldIsEnabled, true)
|
||||
|
||||
if len(params.KnowledgeBaseIDs) > 0 {
|
||||
w.addIn(fieldKnowledgeBaseID, params.KnowledgeBaseIDs)
|
||||
}
|
||||
if len(params.KnowledgeIDs) > 0 {
|
||||
w.addIn(fieldKnowledgeID, params.KnowledgeIDs)
|
||||
}
|
||||
if len(params.TagIDs) > 0 {
|
||||
w.addIn(fieldTagID, params.TagIDs)
|
||||
}
|
||||
if len(params.ExcludeKnowledgeIDs) > 0 {
|
||||
w.addNotIn(fieldKnowledgeID, params.ExcludeKnowledgeIDs)
|
||||
}
|
||||
if len(params.ExcludeChunkIDs) > 0 {
|
||||
w.addNotIn(fieldChunkID, params.ExcludeChunkIDs)
|
||||
}
|
||||
return w
|
||||
}
|
||||
|
||||
// parseEmbeddingLiteral 解析 Doris ARRAY<FLOAT> 通过 MySQL 协议返回的
|
||||
// 字面量字符串(形如 "[1,2,3]")为 []float32。
|
||||
//
|
||||
// CopyIndices 路径需要从源行读出向量本身再写回目标行;此处的解析容错优先:
|
||||
// 不带 [] 也接受、空数组返回 nil。
|
||||
func parseEmbeddingLiteral(raw []byte) ([]float32, error) {
|
||||
s := strings.TrimSpace(string(raw))
|
||||
if s == "" {
|
||||
return nil, nil
|
||||
}
|
||||
s = strings.TrimPrefix(s, "[")
|
||||
s = strings.TrimSuffix(s, "]")
|
||||
if s == "" {
|
||||
return nil, nil
|
||||
}
|
||||
parts := strings.Split(s, ",")
|
||||
out := make([]float32, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
p = strings.TrimSpace(p)
|
||||
if p == "" {
|
||||
continue
|
||||
}
|
||||
f, err := strconv.ParseFloat(p, 32)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, float32(f))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// validateEmbedding 校验向量元素均为有限值。
|
||||
//
|
||||
// strconv.FormatFloat 对 NaN/±Inf 会输出 "NaN"/"+Inf"/"-Inf",
|
||||
// 这些字面量会让 Doris 拼出来的 SQL 报语法错误(或在某些版本下产生未定义结果)。
|
||||
// 上游(嵌入模型)正常情况下不会输出非有限值,但 GPU OOM、上游 bug、
|
||||
// 测试桩都可能触发;这里 fail-fast 比悄悄写脏数据安全。
|
||||
func validateEmbedding(vec []float32) error {
|
||||
for i, v := range vec {
|
||||
if f := float64(v); math.IsNaN(f) || math.IsInf(f, 0) {
|
||||
return errInvalidEmbedding{index: i, value: v}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalizeEmbedding returns a unit-length copy of vec so Doris inner-product
|
||||
// ANN search can preserve cosine-style similarity semantics.
|
||||
func normalizeEmbedding(vec []float32) []float32 {
|
||||
if len(vec) == 0 {
|
||||
return nil
|
||||
}
|
||||
var sumSquares float64
|
||||
for _, value := range vec {
|
||||
f := float64(value)
|
||||
sumSquares += f * f
|
||||
}
|
||||
if sumSquares == 0 {
|
||||
return append([]float32(nil), vec...)
|
||||
}
|
||||
norm := float32(math.Sqrt(sumSquares))
|
||||
normalized := make([]float32, len(vec))
|
||||
for i, value := range vec {
|
||||
normalized[i] = value / norm
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
// errInvalidEmbedding 描述哪个下标含非有限值;用结构体而非 fmt.Errorf
|
||||
// 是为了让上层在日志里能拿到下标做问题定位。
|
||||
type errInvalidEmbedding struct {
|
||||
index int
|
||||
value float32
|
||||
}
|
||||
|
||||
func (e errInvalidEmbedding) Error() string {
|
||||
return "doris: embedding[" + strconv.Itoa(e.index) +
|
||||
"] is not finite: " + strconv.FormatFloat(float64(e.value), 'g', -1, 32)
|
||||
}
|
||||
|
||||
// embeddingLiteral 把 []float32 转为 Doris ARRAY<FLOAT> 字面量字符串:
|
||||
// "[1.23,4.56,...]"。
|
||||
//
|
||||
// 为何不用占位符:go-sql-driver/mysql 不支持 ARRAY 类型的参数绑定,
|
||||
// Doris 端也只接受字面量形式。这里用 strconv.FormatFloat('g' + bitSize=32)
|
||||
// 而不是 fmt.Sprintf("%f", v),原因有二:
|
||||
// 1. fmt 在某些 locale 下会用千分位分隔符,破坏 SQL 语法;
|
||||
// 2. 'g' 比 'f' 短且不会丢精度。
|
||||
//
|
||||
// 注入风险:[]float32 元素是嵌入模型输出的有限位数浮点数,序列化后只可能
|
||||
// 包含 [0-9eE+-.\s] 字符,不会逃逸出字面量上下文。
|
||||
func embeddingLiteral(vec []float32) string {
|
||||
if len(vec) == 0 {
|
||||
return "[]"
|
||||
}
|
||||
var sb strings.Builder
|
||||
sb.Grow(len(vec) * 12)
|
||||
sb.WriteByte('[')
|
||||
for i, v := range vec {
|
||||
if i > 0 {
|
||||
sb.WriteByte(',')
|
||||
}
|
||||
sb.WriteString(strconv.FormatFloat(float64(v), 'g', -1, 32))
|
||||
}
|
||||
sb.WriteByte(']')
|
||||
return sb.String()
|
||||
}
|
||||
@@ -0,0 +1,699 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultTableBaseName = "weknora_embeddings"
|
||||
envDorisTablePrefix = "DORIS_TABLE_PREFIX"
|
||||
)
|
||||
|
||||
// NewDorisRetrieveEngineRepository 创建 Doris 检索引擎仓储。
|
||||
//
|
||||
// 参数:
|
||||
// - db:MySQL 协议的 *sql.DB 实例。调用方负责 SetMaxOpenConns 等参数。
|
||||
// - feHTTPBase:Stream Load 用的 FE HTTP 基地址(含 scheme),例如 "http://doris-fe:8030"。
|
||||
// - username/password:MySQL 与 Stream Load 共用的凭据。
|
||||
// - database:目标数据库名(既用于 MySQL DSN,也用于 Stream Load URL 路径)。
|
||||
// - indexCfg:可空。为 nil 时退化为环境变量 + 默认值(env 路径)。
|
||||
func NewDorisRetrieveEngineRepository(
|
||||
db *sql.DB,
|
||||
feHTTPBase, username, password, database string,
|
||||
indexCfg *types.IndexConfig,
|
||||
) interfaces.RetrieveEngineRepository {
|
||||
log := logger.GetLogger(context.Background())
|
||||
log.Info("[Doris] Initializing Doris retriever engine repository")
|
||||
|
||||
tableBaseName := types.ResolveCollectionName(indexCfg, envDorisTablePrefix, defaultTableBaseName)
|
||||
compatMode, invalidCompatMode := resolveConfiguredDorisCompatMode()
|
||||
if invalidCompatMode != "" {
|
||||
log.Warnf("[Doris] Invalid %s=%q, defaulting to %s", envDorisCompatMode, invalidCompatMode, dorisCompatModeAuto)
|
||||
}
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
httpClient: newDorisStreamLoadHTTPClient(),
|
||||
feHTTPBase: strings.TrimRight(feHTTPBase, "/"),
|
||||
username: username,
|
||||
password: password,
|
||||
database: database,
|
||||
tableBaseName: tableBaseName,
|
||||
bucketsNum: indexCfg.GetBucketsNum(0),
|
||||
replicationNum: indexCfg.GetReplicationNum(0),
|
||||
compatModeRequested: compatMode,
|
||||
}
|
||||
log.Infof("[Doris] Repository initialized: db=%s, base=%s, fe_http=%s, compat_mode=%s",
|
||||
database, tableBaseName, repo.feHTTPBase, repo.compatModeRequested)
|
||||
if os.Getenv(envDorisCompatMode) == "" {
|
||||
log.Infof("[Doris] %s not set, defaulting to %s and probing on first use", envDorisCompatMode, dorisCompatModeAuto)
|
||||
}
|
||||
return repo
|
||||
}
|
||||
|
||||
func (r *dorisRepository) EngineType() types.RetrieverEngineType {
|
||||
return types.DorisRetrieverEngineType
|
||||
}
|
||||
|
||||
func (r *dorisRepository) Support() []types.RetrieverType {
|
||||
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
|
||||
}
|
||||
|
||||
// EstimateStorageSize 估算给定 IndexInfo 列表的存储字节数。
|
||||
//
|
||||
// 参考 Qdrant 的算法:payload 字段长度 + 向量字节 + HNSW 邻居 + 元数据。
|
||||
func (r *dorisRepository) EstimateStorageSize(_ context.Context,
|
||||
indexInfoList []*types.IndexInfo, params map[string]any,
|
||||
) int64 {
|
||||
var total int64
|
||||
for _, info := range indexInfoList {
|
||||
emb := toDorisVectorEmbedding(info, params, dorisCompatModeInnerProductDuplicate)
|
||||
total += calculateStorageSize(emb)
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// Save 写入单条记录到对应维度的表。空向量在 BatchSave 内部统一拒绝。
|
||||
func (r *dorisRepository) Save(ctx context.Context,
|
||||
info *types.IndexInfo, additionalParams map[string]any,
|
||||
) error {
|
||||
return r.BatchSave(ctx, []*types.IndexInfo{info}, additionalParams)
|
||||
}
|
||||
|
||||
// BatchSave 把同一批 IndexInfo 按维度分组;Doris ANN 兼容表使用
|
||||
// DUPLICATE KEY,因此这里显式按 id 做 delete + insert,保持原先的替换语义。
|
||||
func (r *dorisRepository) BatchSave(ctx context.Context,
|
||||
indexInfoList []*types.IndexInfo, additionalParams map[string]any,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(indexInfoList) == 0 {
|
||||
return nil
|
||||
}
|
||||
compatMode, err := r.resolveCompatMode(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
groups := make(map[int][]*DorisVectorEmbedding)
|
||||
for _, info := range indexInfoList {
|
||||
emb := toDorisVectorEmbedding(info, additionalParams, compatMode)
|
||||
if len(emb.Embedding) == 0 {
|
||||
log.Warnf("[Doris] Skipping empty embedding for chunk %s", info.ChunkID)
|
||||
continue
|
||||
}
|
||||
if err := validateEmbedding(emb.Embedding); err != nil {
|
||||
return fmt.Errorf("invalid embedding for chunk %s: %w", info.ChunkID, err)
|
||||
}
|
||||
// 给一个稳定的主键。SourceID 是上层最有意义的"行身份",
|
||||
// 但同 chunk 多 question 的场景下 SourceID 已经唯一,所以直接用它。
|
||||
if emb.ID == "" {
|
||||
emb.ID = emb.SourceID
|
||||
}
|
||||
if emb.ID == "" {
|
||||
emb.ID = uuid.New().String()
|
||||
}
|
||||
dim := len(emb.Embedding)
|
||||
groups[dim] = append(groups[dim], emb)
|
||||
}
|
||||
|
||||
for dim, rows := range groups {
|
||||
if err := r.ensureTable(ctx, dim); err != nil {
|
||||
return err
|
||||
}
|
||||
if compatMode.usesReplaceWrite() {
|
||||
err = r.replaceRows(ctx, r.getTableName(dim), rows)
|
||||
} else {
|
||||
err = r.insertRows(ctx, r.getTableName(dim), rows)
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("batch save dim=%d: %w", dim, err)
|
||||
}
|
||||
log.Infof("[Doris] Saved %d rows to %s", len(rows), r.getTableName(dim))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// insertRows 按列序拼一条多 VALUES 的 INSERT。embedding 列由于
|
||||
// go-sql-driver/mysql 不支持 ARRAY 占位符,必须以字面量形式拼到 SQL 文本中。
|
||||
func (r *dorisRepository) insertRows(ctx context.Context,
|
||||
table string, rows []*DorisVectorEmbedding,
|
||||
) error {
|
||||
if len(rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 9 个普通占位符 + 1 个 embedding 字面量。
|
||||
const perRowPlaceholders = "(?, ?, ?, ?, ?, ?, ?, ?, ?, %s)"
|
||||
|
||||
parts := make([]string, len(rows))
|
||||
args := make([]any, 0, len(rows)*9)
|
||||
for i, e := range rows {
|
||||
parts[i] = fmt.Sprintf(perRowPlaceholders, embeddingLiteral(e.Embedding))
|
||||
args = append(args,
|
||||
e.ID, e.Content, e.SourceID, e.SourceType,
|
||||
e.ChunkID, e.KnowledgeID, e.KnowledgeBaseID, e.TagID,
|
||||
e.IsEnabled,
|
||||
)
|
||||
}
|
||||
|
||||
stmt := fmt.Sprintf("INSERT INTO `%s` (%s) VALUES %s",
|
||||
table,
|
||||
strings.Join(columns, ", "),
|
||||
strings.Join(parts, ", "),
|
||||
)
|
||||
_, err := r.db.ExecContext(ctx, stmt, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
// replaceRows 在 DUPLICATE KEY 表上显式模拟“按 id 覆盖”的语义。
|
||||
func (r *dorisRepository) replaceRows(ctx context.Context,
|
||||
table string, rows []*DorisVectorEmbedding,
|
||||
) error {
|
||||
rows = dedupeRowsByID(rows)
|
||||
if len(rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
ids := make([]string, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
ids = append(ids, row.ID)
|
||||
}
|
||||
if err := r.deleteRowsByID(ctx, table, ids); err != nil {
|
||||
return err
|
||||
}
|
||||
return r.insertRows(ctx, table, rows)
|
||||
}
|
||||
|
||||
func (r *dorisRepository) deleteRowsByID(ctx context.Context,
|
||||
table string, ids []string,
|
||||
) error {
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
placeholders := make([]string, len(ids))
|
||||
args := make([]any, len(ids))
|
||||
for i, id := range ids {
|
||||
placeholders[i] = "?"
|
||||
args[i] = id
|
||||
}
|
||||
|
||||
stmt := fmt.Sprintf("DELETE FROM `%s` WHERE %s IN (%s)",
|
||||
table, fieldID, strings.Join(placeholders, ", "))
|
||||
_, err := r.db.ExecContext(ctx, stmt, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
func dedupeRowsByID(rows []*DorisVectorEmbedding) []*DorisVectorEmbedding {
|
||||
if len(rows) < 2 {
|
||||
return rows
|
||||
}
|
||||
|
||||
out := make([]*DorisVectorEmbedding, 0, len(rows))
|
||||
positions := make(map[string]int, len(rows))
|
||||
for _, row := range rows {
|
||||
if idx, ok := positions[row.ID]; ok {
|
||||
out[idx] = row
|
||||
continue
|
||||
}
|
||||
positions[row.ID] = len(out)
|
||||
out = append(out, row)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// DeleteByChunkIDList 用 chunk_id 列删除。dimension 用于定位具体表。
|
||||
func (r *dorisRepository) DeleteByChunkIDList(ctx context.Context,
|
||||
chunkIDList []string, dimension int, _ string,
|
||||
) error {
|
||||
return r.deleteByField(ctx, fieldChunkID, chunkIDList, dimension)
|
||||
}
|
||||
|
||||
// DeleteByKnowledgeIDList 用 knowledge_id 列删除。
|
||||
func (r *dorisRepository) DeleteByKnowledgeIDList(ctx context.Context,
|
||||
knowledgeIDList []string, dimension int, _ string,
|
||||
) error {
|
||||
return r.deleteByField(ctx, fieldKnowledgeID, knowledgeIDList, dimension)
|
||||
}
|
||||
|
||||
// DeleteBySourceIDList 用 source_id 列删除。
|
||||
func (r *dorisRepository) DeleteBySourceIDList(ctx context.Context,
|
||||
sourceIDList []string, dimension int, _ string,
|
||||
) error {
|
||||
return r.deleteByField(ctx, fieldSourceID, sourceIDList, dimension)
|
||||
}
|
||||
|
||||
// deleteByField 是三个 Delete* 方法的统一实现:
|
||||
// DELETE FROM <table> WHERE <field> IN (?, ?, ...)。
|
||||
func (r *dorisRepository) deleteByField(ctx context.Context,
|
||||
field string, ids []string, dimension int,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
table := r.getTableName(dimension)
|
||||
placeholders := make([]string, len(ids))
|
||||
args := make([]any, len(ids))
|
||||
for i, v := range ids {
|
||||
placeholders[i] = "?"
|
||||
args[i] = v
|
||||
}
|
||||
stmt := fmt.Sprintf("DELETE FROM `%s` WHERE %s IN (%s)",
|
||||
table, field, strings.Join(placeholders, ", "))
|
||||
|
||||
if _, err := r.db.ExecContext(ctx, stmt, args...); err != nil {
|
||||
log.Errorf("[Doris] Delete by %s failed: %v", field, err)
|
||||
return fmt.Errorf("delete by %s: %w", field, err)
|
||||
}
|
||||
log.Infof("[Doris] Deleted %d rows from %s by %s", len(ids), table, field)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Retrieve 根据 RetrieverType 分发到向量检索或关键词检索。
|
||||
func (r *dorisRepository) Retrieve(ctx context.Context,
|
||||
params types.RetrieveParams,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
switch params.RetrieverType {
|
||||
case types.VectorRetrieverType:
|
||||
return r.VectorRetrieve(ctx, params)
|
||||
case types.KeywordsRetrieverType:
|
||||
return r.KeywordsRetrieve(ctx, params)
|
||||
}
|
||||
return nil, fmt.Errorf("invalid retriever type: %v", params.RetrieverType)
|
||||
}
|
||||
|
||||
// VectorRetrieve 对查询向量先做单位化,再调用 inner_product_approximate 做
|
||||
// ANN 搜索;对单位向量而言,inner product 与 cosine similarity 等价,
|
||||
// 因此 score 仍然保持“越大越相似”的语义。
|
||||
func (r *dorisRepository) VectorRetrieve(ctx context.Context,
|
||||
params types.RetrieveParams,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
log := logger.GetLogger(ctx)
|
||||
if err := validateEmbedding(params.Embedding); err != nil {
|
||||
return nil, fmt.Errorf("invalid query embedding: %w", err)
|
||||
}
|
||||
compatMode, err := r.resolveCompatMode(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
queryEmbedding := append([]float32(nil), params.Embedding...)
|
||||
if compatMode.normalizeEmbeddings() {
|
||||
queryEmbedding = normalizeEmbedding(queryEmbedding)
|
||||
}
|
||||
dim := len(params.Embedding)
|
||||
table := r.getTableName(dim)
|
||||
|
||||
exists, err := r.tableExists(ctx, table)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("check table %s: %w", table, err)
|
||||
}
|
||||
if !exists {
|
||||
log.Warnf("[Doris] Table %s does not exist, returning empty results", table)
|
||||
return buildRetrieveResult(nil, types.VectorRetrieverType), nil
|
||||
}
|
||||
|
||||
wb := buildBaseFilter(params)
|
||||
whereClause, whereArgs := wb.build()
|
||||
|
||||
// embedding 必须用字面量,Doris 不支持 LIMIT/OFFSET 使用占位符,必须内联为字面量。
|
||||
// 使用 HAVING 是因为 score 是 SELECT 列别名,WHERE 阶段还看不到。
|
||||
scoreExpr := fmt.Sprintf("inner_product_approximate(`%s`, %s)", fieldEmbedding, embeddingLiteral(queryEmbedding))
|
||||
if compatMode == dorisCompatModeLegacy {
|
||||
scoreExpr = fmt.Sprintf("(1 - cosine_distance_approximate(`%s`, %s))", fieldEmbedding, embeddingLiteral(queryEmbedding))
|
||||
}
|
||||
stmt := fmt.Sprintf(
|
||||
"SELECT %s, %s AS score "+
|
||||
"FROM `%s` WHERE %s "+
|
||||
"HAVING score >= ? "+
|
||||
"ORDER BY score DESC LIMIT %d",
|
||||
strings.Join(columnsForRetrieve, ", "),
|
||||
scoreExpr,
|
||||
table,
|
||||
whereClause,
|
||||
params.TopK,
|
||||
)
|
||||
args := append(whereArgs, params.Threshold)
|
||||
|
||||
rows, err := r.db.QueryContext(ctx, stmt, args...)
|
||||
if err != nil {
|
||||
return nil, r.wrapVectorRetrieveError(table, compatMode, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
results, err := scanRetrieveRows(rows, types.MatchTypeEmbedding)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
log.Infof("[Doris] Vector retrieval found %d results in %s", len(results), table)
|
||||
return buildRetrieveResult(results, types.VectorRetrieverType), nil
|
||||
}
|
||||
|
||||
// KeywordsRetrieve 用 Doris 倒排索引 + MATCH_ANY 做关键词匹配。
|
||||
//
|
||||
// 不需要 jieba 客户端分词:CREATE TABLE 时 idx_content 已经声明了 chinese parser。
|
||||
// 不同维度的表跨表合并取 topK,与 Milvus/Weaviate 现状一致。
|
||||
func (r *dorisRepository) KeywordsRetrieve(ctx context.Context,
|
||||
params types.RetrieveParams,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
log := logger.GetLogger(ctx)
|
||||
query := strings.TrimSpace(params.Query)
|
||||
if query == "" {
|
||||
return buildRetrieveResult(nil, types.KeywordsRetrieverType), nil
|
||||
}
|
||||
|
||||
tables, err := r.listEmbeddingTables(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tables: %w", err)
|
||||
}
|
||||
if len(tables) == 0 {
|
||||
return buildRetrieveResult(nil, types.KeywordsRetrieverType), nil
|
||||
}
|
||||
|
||||
wb := buildBaseFilter(params)
|
||||
whereClause, whereArgs := wb.build()
|
||||
|
||||
var all []*types.IndexWithScore
|
||||
for _, table := range tables {
|
||||
stmt := fmt.Sprintf(
|
||||
"SELECT %s FROM `%s` WHERE %s AND %s MATCH_ANY ? LIMIT %d",
|
||||
strings.Join(columnsForRetrieve, ", "),
|
||||
table, whereClause, fieldContent,
|
||||
params.TopK,
|
||||
)
|
||||
args := append(append([]any{}, whereArgs...), query)
|
||||
|
||||
rows, err := r.db.QueryContext(ctx, stmt, args...)
|
||||
if err != nil {
|
||||
log.Warnf("[Doris] Keyword retrieve in %s failed: %v", table, err)
|
||||
continue
|
||||
}
|
||||
// score 在 KeywordsRetrieve 中固定 1.0,与 Qdrant 行为一致。
|
||||
batch, scanErr := scanRetrieveRows(rows, types.MatchTypeKeywords)
|
||||
_ = rows.Close()
|
||||
if scanErr != nil {
|
||||
return nil, scanErr
|
||||
}
|
||||
all = append(all, batch...)
|
||||
}
|
||||
if len(all) > params.TopK {
|
||||
all = all[:params.TopK]
|
||||
}
|
||||
log.Infof("[Doris] Keywords retrieval found %d results across %d tables", len(all), len(tables))
|
||||
return buildRetrieveResult(all, types.KeywordsRetrieverType), nil
|
||||
}
|
||||
|
||||
// CopyIndices 把源知识库的 chunk 复制到目标知识库,避免重新生成 embedding。
|
||||
//
|
||||
// 与 Qdrant 的实现完全镜像:
|
||||
// - 分页扫描源表
|
||||
// - 按 sourceToTargetChunkIDMap 把 chunk_id 翻译过去
|
||||
// - 处理 source_id 翻译规则(普通 chunk / 生成型问题 / 其他)
|
||||
// - 把目标行写回同一个表
|
||||
func (r *dorisRepository) CopyIndices(ctx context.Context,
|
||||
sourceKnowledgeBaseID string,
|
||||
sourceToTargetKBIDMap map[string]string,
|
||||
sourceToTargetChunkIDMap map[string]string,
|
||||
targetKnowledgeBaseID string,
|
||||
dimension int,
|
||||
_ string,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(sourceToTargetChunkIDMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
if err := r.ensureTable(ctx, dimension); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
table := r.getTableName(dimension)
|
||||
const pageSize = 64
|
||||
offset := 0
|
||||
totalCopied := 0
|
||||
|
||||
for {
|
||||
stmt := fmt.Sprintf(
|
||||
"SELECT %s FROM `%s` WHERE %s = ? ORDER BY %s LIMIT %d OFFSET %d",
|
||||
strings.Join(columnsForCopy, ", "),
|
||||
table, fieldKnowledgeBaseID, fieldID,
|
||||
pageSize, offset,
|
||||
)
|
||||
rows, err := r.db.QueryContext(ctx, stmt, sourceKnowledgeBaseID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("copy indices scan: %w", err)
|
||||
}
|
||||
batch, err := scanCopyRows(rows)
|
||||
_ = rows.Close()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(batch) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
var targets []*DorisVectorEmbedding
|
||||
for _, src := range batch {
|
||||
targetChunkID, ok := sourceToTargetChunkIDMap[src.ChunkID]
|
||||
if !ok {
|
||||
log.Warnf("[Doris] Source chunk %s not in target mapping", src.ChunkID)
|
||||
continue
|
||||
}
|
||||
targetKnowledgeID, ok := sourceToTargetKBIDMap[src.KnowledgeID]
|
||||
if !ok {
|
||||
log.Warnf("[Doris] Source knowledge %s not in target mapping", src.KnowledgeID)
|
||||
continue
|
||||
}
|
||||
|
||||
targetSourceID := translateSourceID(src.SourceID, src.ChunkID, targetChunkID)
|
||||
targets = append(targets, &DorisVectorEmbedding{
|
||||
ID: uuid.New().String(),
|
||||
Content: src.Content,
|
||||
SourceID: targetSourceID,
|
||||
SourceType: src.SourceType,
|
||||
ChunkID: targetChunkID,
|
||||
KnowledgeID: targetKnowledgeID,
|
||||
KnowledgeBaseID: targetKnowledgeBaseID,
|
||||
TagID: src.TagID,
|
||||
IsEnabled: src.IsEnabled,
|
||||
Embedding: src.Embedding,
|
||||
})
|
||||
}
|
||||
|
||||
if len(targets) > 0 {
|
||||
if err := r.insertRows(ctx, table, targets); err != nil {
|
||||
return fmt.Errorf("copy indices insert: %w", err)
|
||||
}
|
||||
totalCopied += len(targets)
|
||||
}
|
||||
|
||||
if len(batch) < pageSize {
|
||||
break
|
||||
}
|
||||
offset += pageSize
|
||||
}
|
||||
log.Infof("[Doris] CopyIndices done, dim=%d, copied=%d", dimension, totalCopied)
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID 实际实现位于 streamload.go,
|
||||
// 会按 compat mode 选择 partial update 或 rewrite rows。
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 私有辅助
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// toDorisVectorEmbedding 把 IndexInfo + 上层传入的 embedding 映射 转换为
|
||||
// Doris 行模型。Embedding 通过 additionalParams[fieldEmbedding] 中的
|
||||
// map[string][]float32 按 SourceID 取出,与 Qdrant/Milvus 完全一致。
|
||||
// inner_product_duplicate 模式会先单位化,legacy 模式保留原始向量。
|
||||
func toDorisVectorEmbedding(
|
||||
info *types.IndexInfo,
|
||||
additionalParams map[string]any,
|
||||
compatMode dorisCompatMode,
|
||||
) *DorisVectorEmbedding {
|
||||
emb := &DorisVectorEmbedding{
|
||||
ID: info.ID,
|
||||
Content: info.Content,
|
||||
SourceID: info.SourceID,
|
||||
SourceType: int(info.SourceType),
|
||||
ChunkID: info.ChunkID,
|
||||
KnowledgeID: info.KnowledgeID,
|
||||
KnowledgeBaseID: info.KnowledgeBaseID,
|
||||
TagID: info.TagID,
|
||||
IsEnabled: info.IsEnabled,
|
||||
}
|
||||
if additionalParams != nil {
|
||||
if v, ok := additionalParams[fieldEmbedding]; ok {
|
||||
if m, ok := v.(map[string][]float32); ok {
|
||||
emb.Embedding = append([]float32(nil), m[info.SourceID]...)
|
||||
if compatMode.normalizeEmbeddings() {
|
||||
emb.Embedding = normalizeEmbedding(emb.Embedding)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return emb
|
||||
}
|
||||
|
||||
func (r *dorisRepository) wrapVectorRetrieveError(table string, compatMode dorisCompatMode, err error) error {
|
||||
if compatMode == dorisCompatModeLegacy {
|
||||
return fmt.Errorf(
|
||||
"vector retrieve %s in Doris compat mode %s: %w. If your Doris build does not support cosine_distance_approximate or ANN on UNIQUE KEY tables, set %s=%s before creating embedding tables. %s is not interchangeable after %s_* tables are created",
|
||||
table,
|
||||
compatMode,
|
||||
err,
|
||||
envDorisCompatMode,
|
||||
dorisCompatModeInnerProductDuplicate,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
)
|
||||
}
|
||||
return fmt.Errorf("vector retrieve %s in Doris compat mode %s: %w", table, compatMode, err)
|
||||
}
|
||||
|
||||
// translateSourceID 把源 SourceID 翻译到目标 SourceID,与 Qdrant 实现完全镜像:
|
||||
// - 普通 chunk:SourceID == ChunkID -> 使用 targetChunkID
|
||||
// - 生成型问题:SourceID == "<chunkID>-<questionID>" -> "<targetChunkID>-<questionID>"
|
||||
// - 其他场景:生成新的 UUID(保持唯一性)
|
||||
func translateSourceID(originalSourceID, sourceChunkID, targetChunkID string) string {
|
||||
switch {
|
||||
case originalSourceID == sourceChunkID:
|
||||
return targetChunkID
|
||||
case strings.HasPrefix(originalSourceID, sourceChunkID+"-"):
|
||||
questionID := strings.TrimPrefix(originalSourceID, sourceChunkID+"-")
|
||||
return fmt.Sprintf("%s-%s", targetChunkID, questionID)
|
||||
default:
|
||||
return uuid.New().String()
|
||||
}
|
||||
}
|
||||
|
||||
// scanRetrieveRows 把 Retrieve 阶段的 rows 反序列化为 IndexWithScore 列表。
|
||||
//
|
||||
// 分两路:
|
||||
// - 列数 == columnsForRetrieve+1:第 N+1 列是 score(向量检索路径)
|
||||
// - 列数 == columnsForRetrieve:score 统一赋 1.0(关键词检索路径)
|
||||
func scanRetrieveRows(rows *sql.Rows, matchType types.MatchType) ([]*types.IndexWithScore, error) {
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
withScore := len(cols) == len(columnsForRetrieve)+1
|
||||
|
||||
var out []*types.IndexWithScore
|
||||
for rows.Next() {
|
||||
var (
|
||||
id, content, sourceID, chunkID string
|
||||
knowledgeID, knowledgeBaseID, tagID string
|
||||
sourceType int
|
||||
isEnabled bool
|
||||
score float64
|
||||
err error
|
||||
)
|
||||
if withScore {
|
||||
err = rows.Scan(&id, &content, &sourceID, &sourceType,
|
||||
&chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled, &score)
|
||||
} else {
|
||||
err = rows.Scan(&id, &content, &sourceID, &sourceType,
|
||||
&chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled)
|
||||
score = 1.0
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("scan row: %w", err)
|
||||
}
|
||||
out = append(out, &types.IndexWithScore{
|
||||
ID: id,
|
||||
Content: content,
|
||||
SourceID: sourceID,
|
||||
SourceType: types.SourceType(sourceType),
|
||||
ChunkID: chunkID,
|
||||
KnowledgeID: knowledgeID,
|
||||
KnowledgeBaseID: knowledgeBaseID,
|
||||
TagID: tagID,
|
||||
Score: score,
|
||||
MatchType: matchType,
|
||||
})
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// scanCopyRows 反序列化 CopyIndices 的分页查询结果。
|
||||
//
|
||||
// 与 scanRetrieveRows 不同,这里需要 embedding 字段(去复制原始向量)。
|
||||
// Doris 的 ARRAY<FLOAT> 通过 mysql 协议返回的是字符串字面量 "[1,2,3]"。
|
||||
func scanCopyRows(rows *sql.Rows) ([]*DorisVectorEmbedding, error) {
|
||||
var out []*DorisVectorEmbedding
|
||||
for rows.Next() {
|
||||
var (
|
||||
id, content, sourceID, chunkID string
|
||||
knowledgeID, knowledgeBaseID, tagID string
|
||||
sourceType int
|
||||
isEnabled bool
|
||||
embeddingRaw sql.RawBytes
|
||||
)
|
||||
if err := rows.Scan(&id, &content, &sourceID, &sourceType,
|
||||
&chunkID, &knowledgeID, &knowledgeBaseID, &tagID, &isEnabled, &embeddingRaw); err != nil {
|
||||
return nil, fmt.Errorf("scan copy row: %w", err)
|
||||
}
|
||||
vec, err := parseEmbeddingLiteral(embeddingRaw)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse embedding: %w", err)
|
||||
}
|
||||
out = append(out, &DorisVectorEmbedding{
|
||||
ID: id,
|
||||
Content: content,
|
||||
SourceID: sourceID,
|
||||
SourceType: sourceType,
|
||||
ChunkID: chunkID,
|
||||
KnowledgeID: knowledgeID,
|
||||
KnowledgeBaseID: knowledgeBaseID,
|
||||
TagID: tagID,
|
||||
IsEnabled: isEnabled,
|
||||
Embedding: vec,
|
||||
})
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// buildRetrieveResult 把 IndexWithScore 列表包装成 RetrieveResult。
|
||||
func buildRetrieveResult(results []*types.IndexWithScore, retrieverType types.RetrieverType) []*types.RetrieveResult {
|
||||
return []*types.RetrieveResult{{
|
||||
Results: results,
|
||||
RetrieverEngineType: types.DorisRetrieverEngineType,
|
||||
RetrieverType: retrieverType,
|
||||
Error: nil,
|
||||
}}
|
||||
}
|
||||
|
||||
// calculateStorageSize 估算单行的存储成本。
|
||||
//
|
||||
// 与 Qdrant 一致:payload 字符串字节 + 向量 (dim*4) + HNSW M*2*8 + 元数据 24。
|
||||
func calculateStorageSize(emb *DorisVectorEmbedding) int64 {
|
||||
var payload int64
|
||||
payload += int64(len(emb.Content))
|
||||
payload += int64(len(emb.SourceID))
|
||||
payload += int64(len(emb.ChunkID))
|
||||
payload += int64(len(emb.KnowledgeID))
|
||||
payload += int64(len(emb.KnowledgeBaseID))
|
||||
payload += int64(len(emb.TagID))
|
||||
payload += 8 // source_type int
|
||||
|
||||
var vec int64
|
||||
var hnsw int64
|
||||
if len(emb.Embedding) > 0 {
|
||||
vec = int64(len(emb.Embedding)) * 4
|
||||
const hnswM = 32
|
||||
hnsw = hnswM * 2 * 8
|
||||
}
|
||||
const metaBytes int64 = 24
|
||||
return payload + vec + hnsw + metaBytes
|
||||
}
|
||||
@@ -0,0 +1,762 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// newTestRepo 构造一个共享 sqlmock 的 dorisRepository,
|
||||
// 默认绕过 ensureTable(initializedTables 已置位),便于专注测试 SQL 形态。
|
||||
//
|
||||
// 返回的 cleanup 用 defer 调用即可。
|
||||
func newTestRepo(t *testing.T) (*dorisRepository, sqlmock.Sqlmock, *httptest.Server, func()) {
|
||||
t.Helper()
|
||||
// httptest binds to loopback, which production SSRF policy correctly blocks.
|
||||
// Make that one test host explicit and restore the process-global policy when
|
||||
// the test completes.
|
||||
secutils.SetSSRFWhitelistFromRaw("127.0.0.1")
|
||||
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
|
||||
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// 默认行为:成功,回 1 行。
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
_ = body
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"Status": "Success",
|
||||
"NumberTotalRows": 1,
|
||||
"NumberLoadedRows": 1,
|
||||
"Label": "test",
|
||||
})
|
||||
}))
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
httpClient: srv.Client(),
|
||||
feHTTPBase: srv.URL,
|
||||
username: "u",
|
||||
password: "p",
|
||||
database: "weknora",
|
||||
tableBaseName: "weknora_embeddings",
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
_ = db.Close()
|
||||
srv.Close()
|
||||
}
|
||||
return repo, mock, srv, cleanup
|
||||
}
|
||||
|
||||
func primeCompatMode(repo *dorisRepository, mode dorisCompatMode, err error) {
|
||||
repo.compatModeResolved = mode
|
||||
repo.compatResolveErr = err
|
||||
repo.compatResolveOnce.Do(func() {})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// query.go:whereBuilder / embeddingLiteral / parseEmbeddingLiteral
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestEmbeddingLiteralRoundTrip(t *testing.T) {
|
||||
t.Run("empty vector returns []", func(t *testing.T) {
|
||||
assert.Equal(t, "[]", embeddingLiteral(nil))
|
||||
assert.Equal(t, "[]", embeddingLiteral([]float32{}))
|
||||
})
|
||||
|
||||
t.Run("does not contain locale-sensitive separators", func(t *testing.T) {
|
||||
s := embeddingLiteral([]float32{1.5, -2.25, 0.001})
|
||||
// 必须只包含数字 / 点 / 负号 / e / 逗号 / 方括号
|
||||
assert.Regexp(t, regexp.MustCompile(`^\[[\-+0-9eE.,]+\]$`), s)
|
||||
})
|
||||
|
||||
t.Run("round trip", func(t *testing.T) {
|
||||
orig := []float32{1.5, -2.25, 0.0625}
|
||||
s := embeddingLiteral(orig)
|
||||
got, err := parseEmbeddingLiteral([]byte(s))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, orig, got)
|
||||
})
|
||||
|
||||
t.Run("parse handles whitespace and missing brackets", func(t *testing.T) {
|
||||
v, err := parseEmbeddingLiteral([]byte(" 1.0 , 2.0 ,3.0 "))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []float32{1.0, 2.0, 3.0}, v)
|
||||
})
|
||||
}
|
||||
|
||||
func TestWhereBuilder(t *testing.T) {
|
||||
t.Run("empty builder returns 1 = 1", func(t *testing.T) {
|
||||
w := &whereBuilder{}
|
||||
clause, args := w.build()
|
||||
assert.Equal(t, "1 = 1", clause)
|
||||
assert.Nil(t, args)
|
||||
})
|
||||
|
||||
t.Run("equal + IN + NOT IN", func(t *testing.T) {
|
||||
w := &whereBuilder{}
|
||||
w.addEqual("is_enabled", true)
|
||||
w.addIn("knowledge_base_id", []string{"kb1", "kb2"})
|
||||
w.addNotIn("chunk_id", []string{"x"})
|
||||
clause, args := w.build()
|
||||
assert.Contains(t, clause, "is_enabled = ?")
|
||||
assert.Contains(t, clause, "knowledge_base_id IN (?, ?)")
|
||||
assert.Contains(t, clause, "chunk_id NOT IN (?)")
|
||||
assert.Equal(t, []any{true, "kb1", "kb2", "x"}, args)
|
||||
})
|
||||
|
||||
t.Run("buildBaseFilter applies all params", func(t *testing.T) {
|
||||
w := buildBaseFilter(types.RetrieveParams{
|
||||
KnowledgeBaseIDs: []string{"kb1"},
|
||||
KnowledgeIDs: []string{"k1", "k2"},
|
||||
TagIDs: []string{"t1"},
|
||||
ExcludeKnowledgeIDs: []string{"k9"},
|
||||
ExcludeChunkIDs: []string{"c9"},
|
||||
})
|
||||
clause, _ := w.build()
|
||||
assert.Contains(t, clause, "is_enabled = ?")
|
||||
assert.Contains(t, clause, "knowledge_base_id IN (?)")
|
||||
assert.Contains(t, clause, "knowledge_id IN (?, ?)")
|
||||
assert.Contains(t, clause, "tag_id IN (?)")
|
||||
assert.Contains(t, clause, "knowledge_id NOT IN (?)")
|
||||
assert.Contains(t, clause, "chunk_id NOT IN (?)")
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// streamload.go:chunkRows / partialUpdateRows
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestChunkRows(t *testing.T) {
|
||||
rows := []map[string]any{
|
||||
{"id": "a", "is_enabled": true},
|
||||
{"id": "b", "is_enabled": false},
|
||||
{"id": "c", "is_enabled": true},
|
||||
}
|
||||
|
||||
t.Run("single batch when fits", func(t *testing.T) {
|
||||
batches := chunkRows(rows, 4096)
|
||||
require.Len(t, batches, 1)
|
||||
assert.Len(t, batches[0], 3)
|
||||
})
|
||||
|
||||
t.Run("splits when exceeding maxBytes", func(t *testing.T) {
|
||||
// 给一个非常小的上限,强制每行单独成段
|
||||
batches := chunkRows(rows, 16)
|
||||
assert.GreaterOrEqual(t, len(batches), 2)
|
||||
|
||||
var total int
|
||||
for _, b := range batches {
|
||||
total += len(b)
|
||||
}
|
||||
assert.Equal(t, len(rows), total)
|
||||
})
|
||||
}
|
||||
|
||||
func TestPartialUpdateRows_HappyPath(t *testing.T) {
|
||||
repo, _, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
// 自定义 server 验证请求形态。
|
||||
var captured *http.Request
|
||||
var capturedBody []byte
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
captured = r
|
||||
capturedBody, _ = io.ReadAll(r.Body)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"Status": "Success",
|
||||
"NumberTotalRows": 2,
|
||||
"NumberLoadedRows": 2,
|
||||
"Label": "ok",
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
repo.feHTTPBase = srv.URL
|
||||
|
||||
rows := []map[string]any{
|
||||
{"id": "id1", "is_enabled": true},
|
||||
{"id": "id2", "is_enabled": false},
|
||||
}
|
||||
require.NoError(t, repo.partialUpdateRows(context.Background(),
|
||||
"weknora_embeddings_768", []string{"id", "is_enabled"}, rows))
|
||||
|
||||
require.NotNil(t, captured)
|
||||
assert.Equal(t, http.MethodPut, captured.Method)
|
||||
assert.Equal(t, "/api/weknora/weknora_embeddings_768/_stream_load", captured.URL.Path)
|
||||
|
||||
assert.Equal(t, "true", captured.Header.Get("partial_columns"))
|
||||
assert.Equal(t, "true", captured.Header.Get("strip_outer_array"))
|
||||
assert.Equal(t, "json", captured.Header.Get("format"))
|
||||
assert.Equal(t, "id,is_enabled", captured.Header.Get("columns"))
|
||||
assert.Equal(t, "APPEND", captured.Header.Get("merge_type"))
|
||||
assert.True(t, strings.HasPrefix(captured.Header.Get("Authorization"), "Basic "))
|
||||
|
||||
var got []map[string]any
|
||||
require.NoError(t, json.Unmarshal(capturedBody, &got))
|
||||
assert.Len(t, got, 2)
|
||||
assert.Equal(t, "id1", got[0]["id"])
|
||||
}
|
||||
|
||||
func TestPartialUpdateRows_FailureSurfaced(t *testing.T) {
|
||||
repo, _, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"Status": "Fail",
|
||||
"Message": "label exists",
|
||||
"ErrorURL": "http://...",
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
repo.feHTTPBase = srv.URL
|
||||
|
||||
err := repo.partialUpdateRows(context.Background(),
|
||||
"t", []string{"id", "is_enabled"},
|
||||
[]map[string]any{{"id": "x", "is_enabled": true}},
|
||||
)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "stream load failed")
|
||||
}
|
||||
|
||||
func TestStreamLoadOnce_BlocksUnsafeTarget(t *testing.T) {
|
||||
repo, _, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
// Override the helper's loopback allowance: link-local metadata endpoints
|
||||
// must be rejected before the HTTP client is called.
|
||||
secutils.SetSSRFWhitelistFromRaw("")
|
||||
repo.feHTTPBase = "http://169.254.169.254"
|
||||
|
||||
err := repo.streamLoadOnce(context.Background(), "t", []string{"id"},
|
||||
[]map[string]any{{"id": "x"}})
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "blocked by SSRF validation")
|
||||
}
|
||||
|
||||
func TestDorisStreamLoadHTTPClient_BlocksUnsafeRedirect(t *testing.T) {
|
||||
secutils.SetSSRFWhitelistFromRaw("127.0.0.1")
|
||||
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, "http://169.254.169.254/latest/meta-data", http.StatusTemporaryRedirect)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
req, err := http.NewRequest(http.MethodPut, server.URL, strings.NewReader("[]"))
|
||||
require.NoError(t, err)
|
||||
_, err = newDorisStreamLoadHTTPClient().Do(req)
|
||||
require.Error(t, err)
|
||||
assert.ErrorIs(t, err, secutils.ErrSSRFRedirectBlocked)
|
||||
}
|
||||
|
||||
func TestDorisStreamLoadHTTPClient_ForwardsAuthorizationToTrustedRedirect(t *testing.T) {
|
||||
secutils.SetSSRFWhitelistFromRaw("127.0.0.1")
|
||||
t.Cleanup(secutils.ResetSSRFWhitelistForTest)
|
||||
|
||||
var gotAuthorization string
|
||||
be := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuthorization = r.Header.Get(headerAuthorization)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer be.Close()
|
||||
|
||||
fe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Redirect(w, r, be.URL, http.StatusTemporaryRedirect)
|
||||
}))
|
||||
defer fe.Close()
|
||||
|
||||
req, err := http.NewRequest(http.MethodPut, fe.URL, strings.NewReader("[]"))
|
||||
require.NoError(t, err)
|
||||
req.Header.Set(headerAuthorization, "Basic trusted-doris-credential")
|
||||
|
||||
resp, err := newDorisStreamLoadHTTPClient().Do(req)
|
||||
require.NoError(t, err)
|
||||
defer resp.Body.Close()
|
||||
assert.Equal(t, "Basic trusted-doris-credential", gotAuthorization)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// repository.go:SQL 形态
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestDeleteByChunkIDList_SQLShape(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
mock.ExpectExec(`DELETE FROM .*weknora_embeddings_768.* WHERE chunk_id IN \(\?, \?\)`).
|
||||
WithArgs("c1", "c2").
|
||||
WillReturnResult(sqlmock.NewResult(0, 2))
|
||||
|
||||
err := repo.DeleteByChunkIDList(context.Background(), []string{"c1", "c2"}, 768, "")
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDeleteByKnowledgeIDList_NoOpOnEmpty(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
require.NoError(t, repo.DeleteByKnowledgeIDList(context.Background(), nil, 768, ""))
|
||||
// 空列表不应触发任何 query
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestVectorRetrieve_SQLShape(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings_3").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(1))
|
||||
|
||||
// Doris does not support placeholders for LIMIT, so TopK is inlined as a
|
||||
// literal and is not a bound argument.
|
||||
mock.ExpectQuery(`SELECT id, content, .*inner_product_approximate.*HAVING score >= \? ORDER BY score DESC LIMIT \d+`).
|
||||
WithArgs(true, 0.5).
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{
|
||||
"id", "content", "source_id", "source_type",
|
||||
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
|
||||
"is_enabled", "score",
|
||||
}).AddRow("id1", "hello", "src", 0, "c1", "k1", "kb1", "t1", true, 0.95),
|
||||
)
|
||||
|
||||
results, err := repo.VectorRetrieve(context.Background(), types.RetrieveParams{
|
||||
Embedding: []float32{1, 2, 3},
|
||||
TopK: 5,
|
||||
Threshold: 0.5,
|
||||
RetrieverType: types.VectorRetrieverType,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
require.Len(t, results[0].Results, 1)
|
||||
assert.Equal(t, "id1", results[0].Results[0].ID)
|
||||
assert.InDelta(t, 0.95, results[0].Results[0].Score, 1e-9)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestVectorRetrieve_SQLShape_LegacyMode(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
primeCompatMode(repo, dorisCompatModeLegacy, nil)
|
||||
|
||||
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings_3").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(1))
|
||||
|
||||
mock.ExpectQuery(`SELECT id, content, .*cosine_distance_approximate.*HAVING score >= \? ORDER BY score DESC LIMIT \d+`).
|
||||
WithArgs(true, 0.5).
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{
|
||||
"id", "content", "source_id", "source_type",
|
||||
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
|
||||
"is_enabled", "score",
|
||||
}).AddRow("id1", "hello", "src", 0, "c1", "k1", "kb1", "t1", true, 0.8),
|
||||
)
|
||||
|
||||
results, err := repo.VectorRetrieve(context.Background(), types.RetrieveParams{
|
||||
Embedding: []float32{1, 2, 3},
|
||||
TopK: 5,
|
||||
Threshold: 0.5,
|
||||
RetrieverType: types.VectorRetrieverType,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
require.Len(t, results[0].Results, 1)
|
||||
assert.Equal(t, "id1", results[0].Results[0].ID)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestNormalizeEmbedding(t *testing.T) {
|
||||
t.Run("returns unit vector copy", func(t *testing.T) {
|
||||
got := normalizeEmbedding([]float32{3, 4})
|
||||
require.Len(t, got, 2)
|
||||
assert.InDelta(t, 0.6, got[0], 1e-6)
|
||||
assert.InDelta(t, 0.8, got[1], 1e-6)
|
||||
})
|
||||
|
||||
t.Run("keeps zero vector finite", func(t *testing.T) {
|
||||
got := normalizeEmbedding([]float32{0, 0})
|
||||
require.Len(t, got, 2)
|
||||
assert.Equal(t, []float32{0, 0}, got)
|
||||
})
|
||||
}
|
||||
|
||||
func TestKeywordsRetrieve_SQLShape(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings\\_%").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).
|
||||
AddRow("weknora_embeddings_768"))
|
||||
|
||||
// LIMIT is inlined (see VectorRetrieve), so only is_enabled and the
|
||||
// MATCH_ANY query string are bound arguments.
|
||||
mock.ExpectQuery(`MATCH_ANY \?`).
|
||||
WithArgs(true, "你好").
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{
|
||||
"id", "content", "source_id", "source_type",
|
||||
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
|
||||
"is_enabled",
|
||||
}).AddRow("id1", "你好世界", "src", 0, "c1", "k1", "kb1", "", true),
|
||||
)
|
||||
|
||||
results, err := repo.KeywordsRetrieve(context.Background(), types.RetrieveParams{
|
||||
Query: "你好",
|
||||
TopK: 3,
|
||||
RetrieverType: types.KeywordsRetrieverType,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
require.Len(t, results[0].Results, 1)
|
||||
assert.Equal(t, "id1", results[0].Results[0].ID)
|
||||
assert.Equal(t, 1.0, results[0].Results[0].Score)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestBatchUpdateChunkEnabledStatus_RewritesRows(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
database: "weknora",
|
||||
tableBaseName: "weknora_embeddings",
|
||||
}
|
||||
primeCompatMode(repo, dorisCompatModeInnerProductDuplicate, nil)
|
||||
|
||||
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings\\_%").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768"))
|
||||
mock.ExpectQuery(`SELECT id, content, source_id, source_type, chunk_id, knowledge_id, knowledge_base_id, tag_id, is_enabled, embedding FROM .*weknora_embeddings_768.* WHERE chunk_id IN`).
|
||||
WithArgs("c1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{
|
||||
"id", "content", "source_id", "source_type",
|
||||
"chunk_id", "knowledge_id", "knowledge_base_id", "tag_id",
|
||||
"is_enabled", "embedding",
|
||||
}).AddRow("row-1", "hello", "src1", 0, "c1", "k1", "kb1", "t1", true, "[1,2,3]"))
|
||||
mock.ExpectExec(`DELETE FROM .*weknora_embeddings_768.* WHERE id IN \(\?\)`).
|
||||
WithArgs("row-1").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectExec(`INSERT INTO .*weknora_embeddings_768.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`).
|
||||
WithArgs("row-1", "hello", "src1", 0, "c1", "k1", "kb1", "t1", false).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
require.NoError(t, repo.BatchUpdateChunkEnabledStatus(
|
||||
context.Background(), map[string]bool{"c1": false}))
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestBatchUpdateChunkEnabledStatus_LegacyModeUsesPartialUpdate(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
primeCompatMode(repo, dorisCompatModeLegacy, nil)
|
||||
|
||||
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings\\_%").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768"))
|
||||
mock.ExpectQuery(`SELECT id, chunk_id FROM .*weknora_embeddings_768.* WHERE chunk_id IN`).
|
||||
WithArgs("c1").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "chunk_id"}).AddRow("row-1", "c1"))
|
||||
|
||||
require.NoError(t, repo.BatchUpdateChunkEnabledStatus(
|
||||
context.Background(), map[string]bool{"c1": false}))
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestEnsureTable_DDLShape(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
database: "weknora",
|
||||
tableBaseName: "weknora_embeddings",
|
||||
bucketsNum: 5,
|
||||
replicationNum: 2,
|
||||
}
|
||||
primeCompatMode(repo, dorisCompatModeInnerProductDuplicate, nil)
|
||||
|
||||
// 表不存在
|
||||
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings_768").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(0))
|
||||
// CREATE TABLE 应包含关键属性和 Doris 支持的 inner_product ANN metric
|
||||
mock.ExpectExec(`CREATE TABLE IF NOT EXISTS .*weknora_embeddings_768.*metric_type"="inner_product".*DUPLICATE KEY\(id\).*BUCKETS 5.*replication_num.*=.*2`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
// SHOW INDEX 一次即返回 ANN 已 FINISHED
|
||||
mock.ExpectQuery(`SHOW INDEX FROM .*weknora_embeddings_768.*`).
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{"Table", "Key_name", "State"}).
|
||||
AddRow("weknora_embeddings_768", "idx_emb", "FINISHED"),
|
||||
)
|
||||
|
||||
require.NoError(t, repo.ensureTable(context.Background(), 768))
|
||||
// waitANNReady 在后台 goroutine 里执行,轮询 ExpectationsWereMet 直到 SHOW INDEX 也被消费。
|
||||
require.Eventually(t, func() bool {
|
||||
return mock.ExpectationsWereMet() == nil
|
||||
}, 2*time.Second, 10*time.Millisecond, "expectations should be met after async ANN poll")
|
||||
}
|
||||
|
||||
func TestEnsureTable_DDLShape_LegacyMode(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
database: "weknora",
|
||||
tableBaseName: "weknora_embeddings",
|
||||
bucketsNum: 5,
|
||||
replicationNum: 2,
|
||||
}
|
||||
primeCompatMode(repo, dorisCompatModeLegacy, nil)
|
||||
|
||||
mock.ExpectQuery(`SELECT COUNT\(1\) FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings_768").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"c"}).AddRow(0))
|
||||
mock.ExpectExec(`CREATE TABLE IF NOT EXISTS .*weknora_embeddings_768.*metric_type"="cosine_distance".*UNIQUE KEY\(id\).*enable_unique_key_merge_on_write.*true`).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectQuery(`SHOW INDEX FROM .*weknora_embeddings_768.*`).
|
||||
WillReturnRows(
|
||||
sqlmock.NewRows([]string{"Table", "Key_name", "State"}).
|
||||
AddRow("weknora_embeddings_768", "idx_emb", "FINISHED"),
|
||||
)
|
||||
|
||||
require.NoError(t, repo.ensureTable(context.Background(), 768))
|
||||
require.Eventually(t, func() bool {
|
||||
return mock.ExpectationsWereMet() == nil
|
||||
}, 2*time.Second, 10*time.Millisecond, "expectations should be met after async ANN poll")
|
||||
}
|
||||
|
||||
func TestBatchSave_SQLShape(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
database: "weknora",
|
||||
tableBaseName: "weknora_embeddings",
|
||||
}
|
||||
primeCompatMode(repo, dorisCompatModeInnerProductDuplicate, nil)
|
||||
repo.initializedTables.Store(3, true) // 跳过 ensureTable
|
||||
|
||||
mock.ExpectExec(`DELETE FROM .*weknora_embeddings_3.* WHERE id IN \(\?\)`).
|
||||
WithArgs("src1").
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectExec(`INSERT INTO .*weknora_embeddings_3.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`).
|
||||
WithArgs(
|
||||
"src1", "hello", "src1", 0,
|
||||
"c1", "k1", "kb1", "",
|
||||
true,
|
||||
).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
err = repo.BatchSave(context.Background(),
|
||||
[]*types.IndexInfo{{
|
||||
Content: "hello",
|
||||
SourceID: "src1",
|
||||
ChunkID: "c1",
|
||||
KnowledgeID: "k1",
|
||||
KnowledgeBaseID: "kb1",
|
||||
IsEnabled: true,
|
||||
}},
|
||||
map[string]any{
|
||||
fieldEmbedding: map[string][]float32{
|
||||
"src1": {1, 2, 3},
|
||||
},
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestBatchSave_SQLShape_LegacyMode(t *testing.T) {
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
repo := &dorisRepository{
|
||||
db: db,
|
||||
database: "weknora",
|
||||
tableBaseName: "weknora_embeddings",
|
||||
}
|
||||
primeCompatMode(repo, dorisCompatModeLegacy, nil)
|
||||
repo.initializedTables.Store(3, true)
|
||||
|
||||
mock.ExpectExec(`INSERT INTO .*weknora_embeddings_3.*VALUES \(\?, \?, \?, \?, \?, \?, \?, \?, \?, \[`).
|
||||
WithArgs(
|
||||
"src1", "hello", "src1", 0,
|
||||
"c1", "k1", "kb1", "",
|
||||
true,
|
||||
).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
err = repo.BatchSave(context.Background(),
|
||||
[]*types.IndexInfo{{
|
||||
Content: "hello",
|
||||
SourceID: "src1",
|
||||
ChunkID: "c1",
|
||||
KnowledgeID: "k1",
|
||||
KnowledgeBaseID: "kb1",
|
||||
IsEnabled: true,
|
||||
}},
|
||||
map[string]any{
|
||||
fieldEmbedding: map[string][]float32{
|
||||
"src1": {1, 2, 3},
|
||||
},
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestResolveCompatMode_RejectsExistingModeSwitch(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
repo.compatModeRequested = dorisCompatModeLegacy
|
||||
|
||||
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings\\_%").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}).AddRow("weknora_embeddings_768"))
|
||||
mock.ExpectQuery(`SHOW CREATE TABLE .*weknora_embeddings_768.*`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"Table", "Create Table"}).AddRow(
|
||||
"weknora_embeddings_768",
|
||||
"CREATE TABLE `weknora_embeddings_768` (id VARCHAR(64)) ENGINE=OLAP DUPLICATE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 10 PROPERTIES(\"replication_num\"=\"1\")",
|
||||
))
|
||||
|
||||
_, err := repo.resolveCompatMode(context.Background())
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "not interchangeable after weknora_embeddings_*")
|
||||
assert.Contains(t, err.Error(), string(dorisCompatModeInnerProductDuplicate))
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestResolveCompatMode_AutoPrefersInnerProductDuplicate(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
repo.compatModeRequested = dorisCompatModeAuto
|
||||
|
||||
mock.ExpectQuery(`SELECT TABLE_NAME FROM information_schema.tables`).
|
||||
WithArgs("weknora", "weknora_embeddings\\_%").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"TABLE_NAME"}))
|
||||
mock.ExpectQuery(`SELECT inner_product_approximate\(\[1.0\],\[1.0\]\)`).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"v"}).AddRow(1.0))
|
||||
mock.ExpectQuery(`SELECT cosine_distance_approximate\(\[1.0\],\[1.0\]\)`).
|
||||
WillReturnError(errors.New("unsupported"))
|
||||
|
||||
mode, err := repo.resolveCompatMode(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, dorisCompatModeInnerProductDuplicate, mode)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Engine wiring
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestEngineTypeAndSupport(t *testing.T) {
|
||||
repo, _, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
assert.Equal(t, types.DorisRetrieverEngineType, repo.EngineType())
|
||||
supports := repo.Support()
|
||||
assert.Contains(t, supports, types.KeywordsRetrieverType)
|
||||
assert.Contains(t, supports, types.VectorRetrieverType)
|
||||
}
|
||||
|
||||
func TestRetrieve_DispatchesByType(t *testing.T) {
|
||||
repo, mock, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
// invalid retriever type -> error,不会触发任何 SQL
|
||||
_, err := repo.Retrieve(context.Background(), types.RetrieveParams{
|
||||
RetrieverType: "unknown",
|
||||
})
|
||||
require.Error(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// regression / utility
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestTranslateSourceID(t *testing.T) {
|
||||
t.Run("plain chunk uses target chunk", func(t *testing.T) {
|
||||
got := translateSourceID("c1", "c1", "tc1")
|
||||
assert.Equal(t, "tc1", got)
|
||||
})
|
||||
t.Run("generated question preserves question id", func(t *testing.T) {
|
||||
got := translateSourceID("c1-q9", "c1", "tc1")
|
||||
assert.Equal(t, "tc1-q9", got)
|
||||
})
|
||||
t.Run("unrecognized source falls back to fresh uuid", func(t *testing.T) {
|
||||
got := translateSourceID("totally-other", "c1", "tc1")
|
||||
assert.NotEqual(t, "totally-other", got)
|
||||
assert.Len(t, got, 36) // UUID 长度
|
||||
})
|
||||
}
|
||||
|
||||
func TestEstimateStorageSize(t *testing.T) {
|
||||
repo, _, _, cleanup := newTestRepo(t)
|
||||
defer cleanup()
|
||||
|
||||
out := repo.EstimateStorageSize(context.Background(),
|
||||
[]*types.IndexInfo{{Content: "hello", ChunkID: "c1", KnowledgeID: "k1", KnowledgeBaseID: "kb1"}},
|
||||
map[string]any{
|
||||
fieldEmbedding: map[string][]float32{
|
||||
"": {1, 2, 3},
|
||||
},
|
||||
},
|
||||
)
|
||||
assert.Greater(t, out, int64(0))
|
||||
}
|
||||
|
||||
// 保证 *sql.Rows 错误不会被吞掉。
|
||||
func TestScanRetrieveRows_Error(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
mock.ExpectQuery("SELECT 1").
|
||||
WillReturnError(errors.New("boom"))
|
||||
|
||||
rows, err := db.Query("SELECT 1")
|
||||
if err != nil {
|
||||
// query 直接失败也算预期
|
||||
assert.Equal(t, "boom", err.Error())
|
||||
return
|
||||
}
|
||||
_, scanErr := scanRetrieveRows(rows, types.MatchTypeEmbedding)
|
||||
if scanErr != nil {
|
||||
assert.Error(t, scanErr)
|
||||
}
|
||||
}
|
||||
|
||||
// silence unused import warning when sql isn't directly used
|
||||
var _ = sql.ErrNoRows
|
||||
@@ -0,0 +1,310 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// 默认的桶数 / 副本数。Doris 在 PROPERTIES 不指定时会用集群默认值,
|
||||
// 这里给一个对单机/小集群更友好的保守值。
|
||||
const (
|
||||
defaultBucketsNum = 10
|
||||
defaultReplicationNum = 1
|
||||
|
||||
// ANN 索引就绪轮询的最大等待时间。索引未就绪不阻塞写入路径,
|
||||
// 只阻塞 ensureTable 自身(首次建表场景),所以 30s 是可接受的。
|
||||
annReadyTimeout = 30 * time.Second
|
||||
annReadyPoll = 1 * time.Second
|
||||
)
|
||||
|
||||
// getTableName 返回某个维度对应的物理表名:<base>_<dim>。
|
||||
//
|
||||
// 与 Qdrant/Milvus/Weaviate 的 collection 命名约定一致,
|
||||
// 这样不同 embedding 模型(不同维度)的数据互不冲突。
|
||||
func (r *dorisRepository) getTableName(dimension int) string {
|
||||
return fmt.Sprintf("%s_%d", r.tableBaseName, dimension)
|
||||
}
|
||||
|
||||
// ensureTable 保证目标维度对应的表已经存在;
|
||||
// 不存在则用 CREATE TABLE IF NOT EXISTS 创建,并在创建后轮询 ANN 索引就绪。
|
||||
//
|
||||
// 该方法在每次 Save / BatchSave 之前调用,结果缓存在 initializedTables 中,
|
||||
// 同一进程内同一 dimension 只会真正打一次 SHOW TABLES + DDL。
|
||||
func (r *dorisRepository) ensureTable(ctx context.Context, dimension int) error {
|
||||
if _, ok := r.initializedTables.Load(dimension); ok {
|
||||
return nil
|
||||
}
|
||||
compatMode, err := r.resolveCompatMode(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log := logger.GetLogger(ctx)
|
||||
tableName := r.getTableName(dimension)
|
||||
|
||||
exists, err := r.tableExists(ctx, tableName)
|
||||
if err != nil {
|
||||
log.Errorf("[Doris] Failed to check table existence: %v", err)
|
||||
return fmt.Errorf("check table existence: %w", err)
|
||||
}
|
||||
|
||||
if !exists {
|
||||
log.Infof("[Doris] Creating table %s with dimension %d in compat mode %s", tableName, dimension, compatMode)
|
||||
if err := r.createTable(ctx, tableName, dimension, compatMode); err != nil {
|
||||
log.Errorf("[Doris] Failed to create table: %v", err)
|
||||
return fmt.Errorf("create table: %w", err)
|
||||
}
|
||||
|
||||
// ANN 索引在 Doris 端异步构建。这里在后台 goroutine 里轮询就绪,
|
||||
// 写入路径不阻塞——索引未就绪期间检索会退化为 brute-force(结果对、速度慢),
|
||||
// 比让首批写入卡 30s 更可接受。
|
||||
go func(tn string) {
|
||||
// 用独立 context(带 timeout),避免请求级 ctx 取消把后台轮询也带走。
|
||||
bgCtx, cancel := context.WithTimeout(context.Background(), annReadyTimeout)
|
||||
defer cancel()
|
||||
if err := r.waitANNReady(bgCtx, tn); err != nil {
|
||||
logger.GetLogger(bgCtx).Warnf(
|
||||
"[Doris] ANN index for %s not ready within %s: %v "+
|
||||
"(queries may fall back to brute force temporarily)",
|
||||
tn, annReadyTimeout, err)
|
||||
return
|
||||
}
|
||||
logger.GetLogger(bgCtx).Infof("[Doris] ANN index for %s ready", tn)
|
||||
}(tableName)
|
||||
}
|
||||
|
||||
r.initializedTables.Store(dimension, true)
|
||||
return nil
|
||||
}
|
||||
|
||||
// tableExists 通过 information_schema 判断表是否存在。
|
||||
//
|
||||
// 不直接用 SHOW TABLES 是因为 Doris 4.1 对 SHOW TABLES LIKE 大小写敏感,
|
||||
// 而 information_schema 与 MySQL 兼容性更好。
|
||||
func (r *dorisRepository) tableExists(ctx context.Context, tableName string) (bool, error) {
|
||||
const q = `SELECT COUNT(1) FROM information_schema.tables
|
||||
WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?`
|
||||
var n int
|
||||
if err := r.db.QueryRowContext(ctx, q, r.database, tableName).Scan(&n); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n > 0, nil
|
||||
}
|
||||
|
||||
// createTable 发出 CREATE TABLE DDL。Doris DDL 是同步的(除 ANN 索引构建外),
|
||||
// 返回成功即代表表已可写。
|
||||
func (r *dorisRepository) createTable(ctx context.Context, tableName string, dimension int, compatMode dorisCompatMode) error {
|
||||
buckets := r.bucketsNum
|
||||
if buckets <= 0 {
|
||||
buckets = defaultBucketsNum
|
||||
}
|
||||
replication := r.replicationNum
|
||||
if replication <= 0 {
|
||||
replication = defaultReplicationNum
|
||||
}
|
||||
|
||||
ddl := buildCreateTableDDL(tableName, dimension, buckets, replication, compatMode)
|
||||
_, err := r.db.ExecContext(ctx, ddl)
|
||||
if err != nil && compatMode == dorisCompatModeLegacy {
|
||||
return fmt.Errorf(
|
||||
"legacy Doris table creation failed: %w. If your Doris build rejects ANN indexes on UNIQUE KEY tables, set %s=%s before creating embedding tables. %s is not interchangeable after %s_* tables are created",
|
||||
err,
|
||||
envDorisCompatMode,
|
||||
dorisCompatModeInnerProductDuplicate,
|
||||
envDorisCompatMode,
|
||||
r.tableBaseName,
|
||||
)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// buildCreateTableDDL 根据维度生成 CREATE TABLE DDL。
|
||||
//
|
||||
// 关键点:
|
||||
// - DUPLICATE KEY(id):兼容当前 Doris/SelectDB 对 ANN 索引的表模型要求。
|
||||
// WeKnora 在 Go 端用 delete + insert 保持按 id 替换的写入语义。
|
||||
// - INVERTED 索引覆盖所有过滤字段 + 中文分词的 content 全文索引。
|
||||
// - ANN 索引使用 HNSW + inner_product;Doris 写入/查询前会对向量单位化,
|
||||
// 因此整体仍保持与其他向量库一致的 cosine 相似度语义。
|
||||
//
|
||||
// 注意:DDL 中 dimension / buckets / replication 三个数值字段是 Go 端格式化拼接的,
|
||||
// 不存在 SQL 注入风险(来源都是受控的 IndexConfig int)。
|
||||
func buildCreateTableDDL(tableName string, dimension, buckets, replication int, compatMode dorisCompatMode) string {
|
||||
metricType := "inner_product"
|
||||
keyMode := "DUPLICATE KEY(id)"
|
||||
properties := fmt.Sprintf("\t\"replication_num\"=\"%d\"", replication)
|
||||
if compatMode == dorisCompatModeLegacy {
|
||||
metricType = "cosine_distance"
|
||||
keyMode = "UNIQUE KEY(id)"
|
||||
properties = fmt.Sprintf("\t\"replication_num\"=\"%d\",\n\t\"enable_unique_key_merge_on_write\"=\"true\"", replication)
|
||||
}
|
||||
|
||||
const tpl = `CREATE TABLE IF NOT EXISTS ` + "`%s`" + ` (
|
||||
id VARCHAR(64) NOT NULL,
|
||||
chunk_id VARCHAR(64),
|
||||
knowledge_id VARCHAR(64),
|
||||
knowledge_base_id VARCHAR(64),
|
||||
source_id VARCHAR(255),
|
||||
source_type INT,
|
||||
tag_id VARCHAR(64),
|
||||
is_enabled BOOLEAN,
|
||||
content TEXT,
|
||||
embedding ARRAY<FLOAT> NOT NULL,
|
||||
INDEX idx_chunk (chunk_id) USING INVERTED,
|
||||
INDEX idx_kb (knowledge_base_id) USING INVERTED,
|
||||
INDEX idx_kid (knowledge_id) USING INVERTED,
|
||||
INDEX idx_src (source_id) USING INVERTED,
|
||||
INDEX idx_tag (tag_id) USING INVERTED,
|
||||
INDEX idx_enabled (is_enabled) USING INVERTED,
|
||||
INDEX idx_content (content) USING INVERTED PROPERTIES("parser"="chinese","support_phrase"="true"),
|
||||
INDEX idx_emb (embedding) USING ANN PROPERTIES(
|
||||
"index_type"="hnsw",
|
||||
"metric_type"="%s",
|
||||
"dim"="%d",
|
||||
"max_degree"="32",
|
||||
"ef_construction"="200"
|
||||
)
|
||||
) ENGINE=OLAP
|
||||
%s
|
||||
DISTRIBUTED BY HASH(id) BUCKETS %d
|
||||
PROPERTIES(
|
||||
%s
|
||||
);`
|
||||
return fmt.Sprintf(tpl, tableName, metricType, dimension, keyMode, buckets, properties)
|
||||
}
|
||||
|
||||
// waitANNReady 轮询 SHOW INDEX,等待 ANN 索引进入 FINISHED 状态。
|
||||
//
|
||||
// Doris 的 ANN 索引在建表后会异步构建,期间查询会退化为 brute-force(结果对,速度慢)。
|
||||
// 此处仅做"尽力而为"的等待:到点未就绪只记 warning,不阻塞写入。
|
||||
func (r *dorisRepository) waitANNReady(ctx context.Context, tableName string) error {
|
||||
deadline := time.Now().Add(annReadyTimeout)
|
||||
for {
|
||||
ready, err := r.annIndexReady(ctx, tableName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if ready {
|
||||
return nil
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
return fmt.Errorf("ann index not ready within %s", annReadyTimeout)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(annReadyPoll):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// annIndexReady 检查 ANN 索引的 State 是否为 FINISHED。
|
||||
//
|
||||
// SHOW INDEX FROM <table> 在 Doris 上返回多列;不同小版本列序略有差异,
|
||||
// 这里以列名匹配(来自 information_schema.statistics + 自定义 view 不可行,
|
||||
// 直接用 SHOW INDEX 然后扫描即可)。
|
||||
//
|
||||
// 兼容策略:如果 SHOW INDEX 返回中找不到 idx_emb 行(极旧版本),视为已就绪,
|
||||
// 避免因为不同 Doris 版本的输出差异把启动卡死。
|
||||
func (r *dorisRepository) annIndexReady(ctx context.Context, tableName string) (bool, error) {
|
||||
rows, err := r.db.QueryContext(ctx,
|
||||
fmt.Sprintf("SHOW INDEX FROM `%s`", tableName))
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
cols, err := rows.Columns()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
keyNameIdx, stateIdx := -1, -1
|
||||
for i, c := range cols {
|
||||
switch strings.ToLower(c) {
|
||||
case "key_name":
|
||||
keyNameIdx = i
|
||||
case "state", "index_state":
|
||||
stateIdx = i
|
||||
}
|
||||
}
|
||||
|
||||
for rows.Next() {
|
||||
// 使用 sql.RawBytes 接收以兼容不同列类型。
|
||||
raw := make([]any, len(cols))
|
||||
ptrs := make([]any, len(cols))
|
||||
for i := range raw {
|
||||
ptrs[i] = &raw[i]
|
||||
}
|
||||
if err := rows.Scan(ptrs...); err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
var keyName, state string
|
||||
if keyNameIdx >= 0 {
|
||||
keyName = bytesToString(raw[keyNameIdx])
|
||||
}
|
||||
if stateIdx >= 0 {
|
||||
state = bytesToString(raw[stateIdx])
|
||||
}
|
||||
|
||||
if keyName != "idx_emb" {
|
||||
continue
|
||||
}
|
||||
if stateIdx < 0 {
|
||||
// 旧版本不暴露 state 列,乐观认为已就绪。
|
||||
return true, nil
|
||||
}
|
||||
if !strings.EqualFold(state, "FINISHED") &&
|
||||
!strings.EqualFold(state, "NORMAL") {
|
||||
return false, nil
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
// 走到这里有两种情况:
|
||||
// 1. 找到了 idx_emb 行,且 state 已是 FINISHED/NORMAL(或 stateIdx<0 的旧版本);
|
||||
// 2. 没找到 idx_emb 行(极旧 Doris 不暴露该索引名);
|
||||
// 都视为已就绪,不阻塞。未就绪的分支已在循环内提前 return false。
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// listEmbeddingTables 返回当前 database 下所有 <base>_% 命名的表,
|
||||
// 用于关键词检索 / 跨维度 BatchUpdate。
|
||||
func (r *dorisRepository) listEmbeddingTables(ctx context.Context) ([]string, error) {
|
||||
const q = `SELECT TABLE_NAME FROM information_schema.tables
|
||||
WHERE TABLE_SCHEMA = ? AND TABLE_NAME LIKE ?`
|
||||
rows, err := r.db.QueryContext(ctx, q, r.database, r.tableBaseName+"\\_%")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var names []string
|
||||
for rows.Next() {
|
||||
var n string
|
||||
if err := rows.Scan(&n); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
names = append(names, n)
|
||||
}
|
||||
return names, rows.Err()
|
||||
}
|
||||
|
||||
// bytesToString 把 SHOW INDEX 返回的 raw any(通常是 []byte 或 string)
|
||||
// 安全转成字符串。
|
||||
func bytesToString(v any) string {
|
||||
switch s := v.(type) {
|
||||
case []byte:
|
||||
return string(s)
|
||||
case string:
|
||||
return s
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
return fmt.Sprintf("%v", s)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,508 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// Stream Load 相关常量。
|
||||
const (
|
||||
// 单批 Stream Load 的 JSON body 上限(保守值,远小于 Doris 默认 streaming_load_max_mb=10240)。
|
||||
// 主要目的是控制单次 HTTP 请求的尾延迟,超过则自动拆批。
|
||||
streamLoadMaxBatchBytes = 1 << 20 // 1 MiB
|
||||
|
||||
// HTTP 头 Authorization 用 Basic auth;Stream Load 也支持 token,
|
||||
// 这里走最常用的用户名/密码方案与 MySQL 协议保持一致。
|
||||
headerAuthorization = "Authorization"
|
||||
headerExpect = "Expect"
|
||||
headerContentType = "Content-Type"
|
||||
)
|
||||
|
||||
// streamLoadResponse 是 Doris FE/BE 返回的 Stream Load 结果体。
|
||||
//
|
||||
// 关键字段:Status 应为 "Success" 或 "Publish Timeout"(后者表示数据已写入但发布事务超时,
|
||||
// 仍视为成功)。其它状态都视为失败。
|
||||
type streamLoadResponse struct {
|
||||
TxnId int64 `json:"TxnId"`
|
||||
Label string `json:"Label"`
|
||||
Status string `json:"Status"`
|
||||
Message string `json:"Message"`
|
||||
NumberTotalRows int64 `json:"NumberTotalRows"`
|
||||
NumberLoadedRows int64 `json:"NumberLoadedRows"`
|
||||
NumberFilteredRows int64 `json:"NumberFilteredRows"`
|
||||
NumberUnselectedRows int64 `json:"NumberUnselectedRows"`
|
||||
LoadBytes int64 `json:"LoadBytes"`
|
||||
LoadTimeMs int64 `json:"LoadTimeMs"`
|
||||
BeginTxnTimeMs int64 `json:"BeginTxnTimeMs"`
|
||||
StreamLoadPutTimeMs int64 `json:"StreamLoadPutTimeMs"`
|
||||
ReadDataTimeMs int64 `json:"ReadDataTimeMs"`
|
||||
WriteDataTimeMs int64 `json:"WriteDataTimeMs"`
|
||||
CommitAndPublishTimeMs int64 `json:"CommitAndPublishTimeMs"`
|
||||
ErrorURL string `json:"ErrorURL"`
|
||||
}
|
||||
|
||||
// streamLoadURL 拼装某张表的 Stream Load HTTP 端点。
|
||||
func (r *dorisRepository) streamLoadURL(table string) string {
|
||||
return fmt.Sprintf("%s/api/%s/%s/_stream_load",
|
||||
r.feHTTPBase, url.PathEscape(r.database), url.PathEscape(table))
|
||||
}
|
||||
|
||||
// newDorisStreamLoadHTTPClient protects both the initial FE request and every
|
||||
// FE -> BE redirect at connection time. Doris requires Basic auth to survive
|
||||
// the 307 redirect, so redirect targets on another host must be explicitly
|
||||
// trusted through the SSRF whitelist before credentials are forwarded.
|
||||
func newDorisStreamLoadHTTPClient() *http.Client {
|
||||
cfg := secutils.DefaultSSRFSafeHTTPClientConfig()
|
||||
// Stream Load calls carry their own context deadline. Preserve the previous
|
||||
// client behaviour instead of imposing the generic 30-second HTTP timeout.
|
||||
cfg.Timeout = 0
|
||||
|
||||
client := secutils.NewSSRFSafeHTTPClient(cfg)
|
||||
ssrfCheckRedirect := client.CheckRedirect
|
||||
client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
|
||||
// Keep the shared client's redirect limit, scheme/target validation and
|
||||
// sensitive-header stripping as the authoritative SSRF behaviour.
|
||||
if err := ssrfCheckRedirect(req, via); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(via) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
sourceHost := via[len(via)-1].URL.Hostname()
|
||||
targetHost := req.URL.Hostname()
|
||||
if !strings.EqualFold(sourceHost, targetHost) && !secutils.IsSSRFWhitelisted(targetHost) {
|
||||
return fmt.Errorf(
|
||||
"stream load redirect blocked: target host %q is not trusted to receive credentials",
|
||||
targetHost,
|
||||
)
|
||||
}
|
||||
|
||||
// The shared client deliberately strips Authorization on cross-host
|
||||
// redirects. Doris is the exceptional case where an explicitly trusted
|
||||
// BE needs the same Basic credentials after the FE's 307 response.
|
||||
req.Header.Set(headerAuthorization, via[0].Header.Get(headerAuthorization))
|
||||
return nil
|
||||
}
|
||||
return client
|
||||
}
|
||||
|
||||
// partialUpdateRows 把若干行通过 Stream Load 的 partial update 模式写回目标表。
|
||||
//
|
||||
// columns 是参与本次 partial update 的列(必须包含 UNIQUE KEY 列,即 "id")。
|
||||
// rows 中每一项是一个与 columns 等长的字段值数组。
|
||||
//
|
||||
// 实现要点:
|
||||
// 1. 用 JSON 数组的 body 形式,header 加 strip_outer_array=true。
|
||||
// 2. 设置 partial_columns=true、merge_type=APPEND,触发 Doris 的 partial update 模式
|
||||
// (Doris 4.1 + UNIQUE KEY MoW 表的标准玩法)。
|
||||
// 3. 按 streamLoadMaxBatchBytes 自动拆批,避免单次过大。
|
||||
// 4. 处理 307:Doris 的 FE 会 redirect 到 BE,net/http 默认会跟随;
|
||||
// 这里需要保证 GetBody 可重发(已通过 bytes.NewReader 构造 Body 满足)。
|
||||
func (r *dorisRepository) partialUpdateRows(ctx context.Context,
|
||||
table string, columns []string, rows []map[string]any,
|
||||
) error {
|
||||
if len(rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
for _, batch := range chunkRows(rows, streamLoadMaxBatchBytes) {
|
||||
if err := r.streamLoadOnce(ctx, table, columns, batch); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// streamLoadOnce 发出一次 Stream Load HTTP 请求。
|
||||
func (r *dorisRepository) streamLoadOnce(ctx context.Context,
|
||||
table string, columns []string, rows []map[string]any,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
body, err := json.Marshal(rows)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal stream load body: %w", err)
|
||||
}
|
||||
|
||||
url := r.streamLoadURL(table)
|
||||
// Validate at the final outbound boundary as well as when user-supplied
|
||||
// vector-store configuration is created. This covers stored/env configs and
|
||||
// keeps the tainted value from reaching http.Client.Do unchecked.
|
||||
if err := secutils.ValidateURLForSSRF(url); err != nil {
|
||||
return fmt.Errorf("stream load URL blocked by SSRF validation: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return fmt.Errorf("build stream load request: %w", err)
|
||||
}
|
||||
|
||||
// GetBody 让 redirect 时可以重新读 body(FE -> BE 的 307 需要重发 PUT body)。
|
||||
bodyCopy := append([]byte(nil), body...)
|
||||
req.GetBody = func() (io.ReadCloser, error) {
|
||||
return io.NopCloser(bytes.NewReader(bodyCopy)), nil
|
||||
}
|
||||
req.ContentLength = int64(len(body))
|
||||
|
||||
auth := base64.StdEncoding.EncodeToString([]byte(r.username + ":" + r.password))
|
||||
req.Header.Set(headerAuthorization, "Basic "+auth)
|
||||
req.Header.Set(headerExpect, "100-continue")
|
||||
req.Header.Set(headerContentType, "application/json")
|
||||
req.Header.Set("format", "json")
|
||||
req.Header.Set("strip_outer_array", "true")
|
||||
req.Header.Set("partial_columns", "true")
|
||||
req.Header.Set("columns", strings.Join(columns, ","))
|
||||
req.Header.Set("merge_type", "APPEND")
|
||||
|
||||
resp, err := r.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("stream load HTTP: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read stream load response: %w", err)
|
||||
}
|
||||
if resp.StatusCode/100 != 2 {
|
||||
return fmt.Errorf("stream load HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var result streamLoadResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return fmt.Errorf("decode stream load response: %w (raw=%s)", err, string(respBody))
|
||||
}
|
||||
|
||||
switch result.Status {
|
||||
case "Success", "Publish Timeout":
|
||||
log.Infof("[Doris] Stream load %s OK: rows=%d, loaded=%d, label=%s",
|
||||
table, result.NumberTotalRows, result.NumberLoadedRows, result.Label)
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("stream load failed: status=%s msg=%s err_url=%s",
|
||||
result.Status, result.Message, result.ErrorURL)
|
||||
}
|
||||
}
|
||||
|
||||
// chunkRows 把行按累积 JSON 体大小切分,每段不超过 maxBytes。
|
||||
//
|
||||
// 注意:JSON 序列化的实际开销大约是 marshal 后的字节数,而单行 marshal
|
||||
// 加上逗号 + 数组括号约等于本估算。这里用粗略估计避免每段都 marshal。
|
||||
func chunkRows(rows []map[string]any, maxBytes int) [][]map[string]any {
|
||||
if len(rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var (
|
||||
out [][]map[string]any
|
||||
curr []map[string]any
|
||||
size int
|
||||
header = 2 // "[" + "]"
|
||||
)
|
||||
|
||||
for _, row := range rows {
|
||||
raw, err := json.Marshal(row)
|
||||
if err != nil {
|
||||
// marshal 失败时把这一行单独成段,由上层 streamLoadOnce 再次 marshal 报错。
|
||||
if len(curr) > 0 {
|
||||
out = append(out, curr)
|
||||
}
|
||||
out = append(out, []map[string]any{row})
|
||||
curr = nil
|
||||
size = 0
|
||||
continue
|
||||
}
|
||||
// 加上逗号位(除第一行之外)。
|
||||
need := len(raw)
|
||||
if len(curr) > 0 {
|
||||
need++
|
||||
}
|
||||
if size+need+header > maxBytes && len(curr) > 0 {
|
||||
out = append(out, curr)
|
||||
curr = nil
|
||||
size = 0
|
||||
}
|
||||
curr = append(curr, row)
|
||||
size += need
|
||||
}
|
||||
if len(curr) > 0 {
|
||||
out = append(out, curr)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 业务面方法:BatchUpdateChunkEnabledStatus / BatchUpdateChunkTagID
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// BatchUpdateChunkEnabledStatus 批量更新 chunk 的 is_enabled 字段。
|
||||
// legacy 模式走 Stream Load partial update;inner_product_duplicate 模式改为读整行后 replaceRows 写回。
|
||||
func (r *dorisRepository) BatchUpdateChunkEnabledStatus(ctx context.Context,
|
||||
chunkStatusMap map[string]bool,
|
||||
) error {
|
||||
if len(chunkStatusMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
compatMode, err := r.resolveCompatMode(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !compatMode.usesRewriteChunkUpdates() {
|
||||
return r.batchUpdateChunkEnabledStatusLegacy(ctx, chunkStatusMap)
|
||||
}
|
||||
|
||||
chunkIDs := make([]string, 0, len(chunkStatusMap))
|
||||
for id := range chunkStatusMap {
|
||||
chunkIDs = append(chunkIDs, id)
|
||||
}
|
||||
|
||||
return r.rewriteChunkRows(ctx, chunkIDs, func(row *DorisVectorEmbedding) bool {
|
||||
enabled, ok := chunkStatusMap[row.ChunkID]
|
||||
if !ok || row.IsEnabled == enabled {
|
||||
return false
|
||||
}
|
||||
row.IsEnabled = enabled
|
||||
return true
|
||||
}, "rewrite is_enabled")
|
||||
}
|
||||
|
||||
// BatchUpdateChunkTagID 批量更新 chunk 的 tag_id 字段。逻辑与 EnabledStatus 一致。
|
||||
func (r *dorisRepository) BatchUpdateChunkTagID(ctx context.Context,
|
||||
chunkTagMap map[string]string,
|
||||
) error {
|
||||
if len(chunkTagMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
compatMode, err := r.resolveCompatMode(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !compatMode.usesRewriteChunkUpdates() {
|
||||
return r.batchUpdateChunkTagIDLegacy(ctx, chunkTagMap)
|
||||
}
|
||||
|
||||
chunkIDs := make([]string, 0, len(chunkTagMap))
|
||||
for id := range chunkTagMap {
|
||||
chunkIDs = append(chunkIDs, id)
|
||||
}
|
||||
|
||||
return r.rewriteChunkRows(ctx, chunkIDs, func(row *DorisVectorEmbedding) bool {
|
||||
tagID, ok := chunkTagMap[row.ChunkID]
|
||||
if !ok || row.TagID == tagID {
|
||||
return false
|
||||
}
|
||||
row.TagID = tagID
|
||||
return true
|
||||
}, "rewrite tag_id")
|
||||
}
|
||||
|
||||
func (r *dorisRepository) rewriteChunkRows(ctx context.Context,
|
||||
chunkIDs []string,
|
||||
mutate func(*DorisVectorEmbedding) bool,
|
||||
action string,
|
||||
) error {
|
||||
if len(chunkIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
tables, err := r.listEmbeddingTables(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("list tables: %w", err)
|
||||
}
|
||||
|
||||
for _, table := range tables {
|
||||
rows, err := r.loadRowsByChunkIDs(ctx, table, chunkIDs)
|
||||
if err != nil {
|
||||
return fmt.Errorf("load chunk rows from %s: %w", table, err)
|
||||
}
|
||||
|
||||
updated := make([]*DorisVectorEmbedding, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
if !mutate(row) {
|
||||
continue
|
||||
}
|
||||
updated = append(updated, row)
|
||||
}
|
||||
if len(updated) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
if err := r.replaceRows(ctx, table, updated); err != nil {
|
||||
return fmt.Errorf("%s in %s: %w", action, table, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *dorisRepository) batchUpdateChunkEnabledStatusLegacy(ctx context.Context,
|
||||
chunkStatusMap map[string]bool,
|
||||
) error {
|
||||
chunkIDs := make([]string, 0, len(chunkStatusMap))
|
||||
for id := range chunkStatusMap {
|
||||
chunkIDs = append(chunkIDs, id)
|
||||
}
|
||||
|
||||
mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
byTable := make(map[string][]map[string]any)
|
||||
for chunkID, locations := range mapping {
|
||||
enabled, ok := chunkStatusMap[chunkID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, loc := range locations {
|
||||
byTable[loc.table] = append(byTable[loc.table], map[string]any{
|
||||
fieldID: loc.id,
|
||||
fieldIsEnabled: enabled,
|
||||
})
|
||||
}
|
||||
}
|
||||
for table, rows := range byTable {
|
||||
if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldIsEnabled}, rows); err != nil {
|
||||
return fmt.Errorf("partial update is_enabled in %s: %w", table, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *dorisRepository) batchUpdateChunkTagIDLegacy(ctx context.Context,
|
||||
chunkTagMap map[string]string,
|
||||
) error {
|
||||
chunkIDs := make([]string, 0, len(chunkTagMap))
|
||||
for id := range chunkTagMap {
|
||||
chunkIDs = append(chunkIDs, id)
|
||||
}
|
||||
|
||||
mapping, err := r.lookupChunkRowKeys(ctx, chunkIDs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
byTable := make(map[string][]map[string]any)
|
||||
for chunkID, locations := range mapping {
|
||||
tagID, ok := chunkTagMap[chunkID]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, loc := range locations {
|
||||
byTable[loc.table] = append(byTable[loc.table], map[string]any{
|
||||
fieldID: loc.id,
|
||||
fieldTagID: tagID,
|
||||
})
|
||||
}
|
||||
}
|
||||
for table, rows := range byTable {
|
||||
if err := r.partialUpdateRows(ctx, table, []string{fieldID, fieldTagID}, rows); err != nil {
|
||||
return fmt.Errorf("partial update tag_id in %s: %w", table, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *dorisRepository) loadRowsByChunkIDs(ctx context.Context,
|
||||
table string, chunkIDs []string,
|
||||
) ([]*DorisVectorEmbedding, error) {
|
||||
if len(chunkIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
placeholders := make([]string, len(chunkIDs))
|
||||
args := make([]any, len(chunkIDs))
|
||||
for i, v := range chunkIDs {
|
||||
placeholders[i] = "?"
|
||||
args[i] = v
|
||||
}
|
||||
|
||||
stmt := fmt.Sprintf(
|
||||
"SELECT %s FROM `%s` WHERE %s IN (%s)",
|
||||
strings.Join(columnsForCopy, ", "),
|
||||
table,
|
||||
fieldChunkID,
|
||||
strings.Join(placeholders, ", "),
|
||||
)
|
||||
rows, err := r.db.QueryContext(ctx, stmt, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
batch, err := scanCopyRows(rows)
|
||||
_ = rows.Close()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return batch, nil
|
||||
}
|
||||
|
||||
// rowLocation 表示某行在哪个表里、主键 id 是什么。
|
||||
type rowLocation struct {
|
||||
table string
|
||||
id string
|
||||
}
|
||||
|
||||
// lookupChunkRowKeys 查询给定的 chunkIDs 在所有 <base>_<dim> 表中的物理位置:
|
||||
// - key:chunk_id
|
||||
// - value:[(table, id), ...],因为同一 chunk 可能在多个维度的表里都有副本。
|
||||
//
|
||||
// 跨表查询使用 listEmbeddingTables 列出的所有匹配表;每张表执行一次
|
||||
// SELECT id, chunk_id FROM <table> WHERE chunk_id IN (?, ?, ...)。
|
||||
func (r *dorisRepository) lookupChunkRowKeys(ctx context.Context,
|
||||
chunkIDs []string,
|
||||
) (map[string][]rowLocation, error) {
|
||||
if len(chunkIDs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
tables, err := r.listEmbeddingTables(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tables: %w", err)
|
||||
}
|
||||
if len(tables) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
placeholders := make([]string, len(chunkIDs))
|
||||
args := make([]any, len(chunkIDs))
|
||||
for i, v := range chunkIDs {
|
||||
placeholders[i] = "?"
|
||||
args[i] = v
|
||||
}
|
||||
|
||||
out := make(map[string][]rowLocation)
|
||||
for _, table := range tables {
|
||||
stmt := fmt.Sprintf(
|
||||
"SELECT %s, %s FROM `%s` WHERE %s IN (%s)",
|
||||
fieldID, fieldChunkID, table, fieldChunkID, strings.Join(placeholders, ", "),
|
||||
)
|
||||
rows, err := r.db.QueryContext(ctx, stmt, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("lookup chunk row keys in %s: %w", table, err)
|
||||
}
|
||||
for rows.Next() {
|
||||
var id, chunkID string
|
||||
if err := rows.Scan(&id, &chunkID); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, fmt.Errorf("scan row keys: %w", err)
|
||||
}
|
||||
out[chunkID] = append(out[chunkID], rowLocation{table: table, id: id})
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
_ = rows.Close()
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
package doris
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"net/http"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// dorisRepository 是 Apache Doris 4.1 的检索引擎仓储实现。
|
||||
//
|
||||
// 通信通道:
|
||||
// - 读写主链路:MySQL 协议(database/sql + go-sql-driver/mysql),FE 默认 9030 端口。
|
||||
// - Stream Load:HTTP(FE 默认 8030 端口),legacy 模式用于 partial update。
|
||||
//
|
||||
// 表结构按维度分表:<tableBaseName>_<dim>。
|
||||
// 兼容模式由 DORIS_COMPAT_MODE 决定:
|
||||
// - legacy:UNIQUE KEY(id) + cosine_distance ANN + Stream Load partial update
|
||||
// - inner_product_duplicate:DUPLICATE KEY(id) + normalized inner product + delete/insert rewrite
|
||||
//
|
||||
// 该设置在 embedding 表创建后不可直接互换;切换模式前需要重建这些表。
|
||||
//
|
||||
// 与 Qdrant/Milvus/Weaviate 一样,initializedTables 缓存"已确保存在"的维度,
|
||||
// 避免每次写入都打 SHOW TABLES。
|
||||
type dorisRepository struct {
|
||||
db *sql.DB
|
||||
|
||||
httpClient *http.Client
|
||||
// fe HTTP base,例如 "http://doris-fe:8030"。Stream Load 路径
|
||||
// 由 streamLoadURL(table) 拼接:<feHTTPBase>/api/<database>/<table>/_stream_load。
|
||||
feHTTPBase string
|
||||
|
||||
username string
|
||||
password string
|
||||
database string
|
||||
|
||||
tableBaseName string
|
||||
bucketsNum int // 0 -> default 10
|
||||
replicationNum int // 0 -> default 1
|
||||
compatModeRequested dorisCompatMode
|
||||
compatModeResolved dorisCompatMode
|
||||
compatResolveOnce sync.Once
|
||||
compatResolveErr error
|
||||
|
||||
// 已经确保过 ensureTable 的维度集合:dim -> true。
|
||||
initializedTables sync.Map
|
||||
}
|
||||
|
||||
// DorisVectorEmbedding 是落到 Doris 表里的一行的领域模型。
|
||||
//
|
||||
// 字段顺序与 schema.go 中的 INSERT 列序保持一致,
|
||||
// 调整时需要同时更新 createInsert 与 columns。
|
||||
type DorisVectorEmbedding struct {
|
||||
ID string
|
||||
Content string
|
||||
SourceID string
|
||||
SourceType int
|
||||
ChunkID string
|
||||
KnowledgeID string
|
||||
KnowledgeBaseID string
|
||||
TagID string
|
||||
IsEnabled bool
|
||||
Embedding []float32
|
||||
}
|
||||
|
||||
// DorisVectorEmbeddingWithScore 是检索结果的领域模型,
|
||||
// Score 在向量检索时按当前 compat mode 计算,在关键词检索时统一赋 1.0。
|
||||
type DorisVectorEmbeddingWithScore struct {
|
||||
DorisVectorEmbedding
|
||||
Score float64
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package elasticsearch
|
||||
|
||||
import (
|
||||
"maps"
|
||||
"slices"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// VectorEmbedding defines the Elasticsearch document structure for vector embeddings
|
||||
type VectorEmbedding struct {
|
||||
Content string `json:"content" gorm:"column:content;not null"` // Text content of the chunk
|
||||
SourceID string `json:"source_id" gorm:"column:source_id;not null"` // ID of the source document
|
||||
SourceType int `json:"source_type" gorm:"column:source_type;not null"` // Type of the source document
|
||||
ChunkID string `json:"chunk_id" gorm:"column:chunk_id"` // Unique ID of the text chunk
|
||||
KnowledgeID string `json:"knowledge_id" gorm:"column:knowledge_id"` // ID of the knowledge item
|
||||
KnowledgeBaseID string `json:"knowledge_base_id" gorm:"column:knowledge_base_id"` // ID of the knowledge base
|
||||
TagID string `json:"tag_id" gorm:"column:tag_id"` // Tag ID for categorization
|
||||
Embedding []float32 `json:"embedding" gorm:"column:embedding;not null"` // Vector embedding of the content
|
||||
IsEnabled bool `json:"is_enabled"` // Whether the chunk is enabled
|
||||
IsRecommended bool `json:"is_recommended"` // Whether the chunk is recommended
|
||||
}
|
||||
|
||||
// VectorEmbeddingWithScore extends VectorEmbedding with similarity score
|
||||
type VectorEmbeddingWithScore struct {
|
||||
VectorEmbedding
|
||||
Score float64 // Similarity score from vector search
|
||||
}
|
||||
|
||||
// ToDBVectorEmbedding converts IndexInfo to Elasticsearch document format
|
||||
func ToDBVectorEmbedding(embedding *types.IndexInfo, additionalParams map[string]interface{}) *VectorEmbedding {
|
||||
vector := &VectorEmbedding{
|
||||
Content: embedding.Content,
|
||||
SourceID: embedding.SourceID,
|
||||
SourceType: int(embedding.SourceType),
|
||||
ChunkID: embedding.ChunkID,
|
||||
KnowledgeID: embedding.KnowledgeID,
|
||||
KnowledgeBaseID: embedding.KnowledgeBaseID,
|
||||
TagID: embedding.TagID,
|
||||
IsEnabled: embedding.IsEnabled,
|
||||
IsRecommended: embedding.IsRecommended,
|
||||
}
|
||||
// Add embedding data if available in additionalParams
|
||||
if additionalParams != nil && slices.Contains(slices.Collect(maps.Keys(additionalParams)), "embedding") {
|
||||
if embeddingMap, ok := additionalParams["embedding"].(map[string][]float32); ok {
|
||||
vector.Embedding = embeddingMap[embedding.SourceID]
|
||||
}
|
||||
}
|
||||
// Get is_enabled from additionalParams if available
|
||||
if additionalParams != nil {
|
||||
if chunkEnabledMap, ok := additionalParams["chunk_enabled"].(map[string]bool); ok {
|
||||
if enabled, exists := chunkEnabledMap[embedding.ChunkID]; exists {
|
||||
vector.IsEnabled = enabled
|
||||
}
|
||||
}
|
||||
}
|
||||
return vector
|
||||
}
|
||||
|
||||
// FromDBVectorEmbeddingWithScore converts Elasticsearch document to IndexWithScore domain model
|
||||
func FromDBVectorEmbeddingWithScore(id string,
|
||||
embedding *VectorEmbeddingWithScore,
|
||||
matchType types.MatchType,
|
||||
) *types.IndexWithScore {
|
||||
return &types.IndexWithScore{
|
||||
ID: id,
|
||||
SourceID: embedding.SourceID,
|
||||
SourceType: types.SourceType(embedding.SourceType),
|
||||
ChunkID: embedding.ChunkID,
|
||||
KnowledgeID: embedding.KnowledgeID,
|
||||
KnowledgeBaseID: embedding.KnowledgeBaseID,
|
||||
TagID: embedding.TagID,
|
||||
Content: embedding.Content,
|
||||
Score: embedding.Score,
|
||||
MatchType: matchType,
|
||||
IsEnabled: embedding.IsEnabled,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,820 @@
|
||||
package v8
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
elasticsearchRetriever "github.com/Tencent/WeKnora/internal/application/repository/retriever/elasticsearch"
|
||||
"github.com/Tencent/WeKnora/internal/config"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
typesLocal "github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/elastic/go-elasticsearch/v8"
|
||||
"github.com/elastic/go-elasticsearch/v8/typedapi/core/search"
|
||||
"github.com/elastic/go-elasticsearch/v8/typedapi/types"
|
||||
"github.com/elastic/go-elasticsearch/v8/typedapi/types/enums/scriptlanguage"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// elasticsearchRepository implements the RetrieveEngineRepository interface for Elasticsearch v8
|
||||
type elasticsearchRepository struct {
|
||||
client *elasticsearch.TypedClient // Elasticsearch client instance
|
||||
index string // Name of the Elasticsearch index to use
|
||||
useKeywordSuffix bool // Whether to append .keyword suffix to ID field names in queries
|
||||
numberOfShards int // Shard count for index creation (0 = ES default)
|
||||
numberOfReplicas int // Replica count for index creation (-1 = unset, use ES default)
|
||||
}
|
||||
|
||||
// NewElasticsearchEngineRepository creates and initializes a new Elasticsearch v8 repository.
|
||||
// indexCfg is optional — pass nil to use env var / default values (env path).
|
||||
func NewElasticsearchEngineRepository(client *elasticsearch.TypedClient,
|
||||
config *config.Config,
|
||||
indexCfg *typesLocal.IndexConfig,
|
||||
) interfaces.RetrieveEngineRepository {
|
||||
log := logger.GetLogger(context.Background())
|
||||
log.Info("[Elasticsearch] Initializing Elasticsearch v8 retriever engine repository")
|
||||
|
||||
indexName := typesLocal.ResolveIndexName(indexCfg, "ELASTICSEARCH_INDEX", "xwrag_default")
|
||||
|
||||
// Create repository instance and ensure index exists
|
||||
res := &elasticsearchRepository{
|
||||
client: client,
|
||||
index: indexName,
|
||||
numberOfShards: indexCfg.GetNumberOfShards(0),
|
||||
numberOfReplicas: indexCfg.GetNumberOfReplicas(-1),
|
||||
}
|
||||
if err := res.createIndexIfNotExists(context.Background()); err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to create index: %v", err)
|
||||
} else {
|
||||
log.Info("[Elasticsearch] Successfully initialized repository")
|
||||
}
|
||||
res.detectFieldTypes(context.Background())
|
||||
return res
|
||||
}
|
||||
|
||||
// idField returns the query field name for an ID field, appending ".keyword"
|
||||
// suffix when the index uses text-type mappings with keyword sub-fields.
|
||||
func (e *elasticsearchRepository) idField(name string) string {
|
||||
if e.useKeywordSuffix {
|
||||
return name + ".keyword"
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
// detectFieldTypes inspects the index mapping to determine whether ID fields
|
||||
// are mapped as "keyword" (no suffix needed) or "text" (needs ".keyword" suffix).
|
||||
func (e *elasticsearchRepository) detectFieldTypes(ctx context.Context) {
|
||||
log := logger.GetLogger(ctx)
|
||||
|
||||
mappingResp, err := e.client.Indices.GetMapping().Index(e.index).Do(ctx)
|
||||
if err != nil {
|
||||
log.Warnf("[Elasticsearch] Failed to get index mapping, defaulting to .keyword suffix: %v", err)
|
||||
e.useKeywordSuffix = true
|
||||
return
|
||||
}
|
||||
|
||||
indexMapping, ok := mappingResp[e.index]
|
||||
if !ok {
|
||||
log.Warnf("[Elasticsearch] Index %s not found in mapping response, defaulting to .keyword suffix", e.index)
|
||||
e.useKeywordSuffix = true
|
||||
return
|
||||
}
|
||||
|
||||
if prop, ok := indexMapping.Mappings.Properties["chunk_id"]; ok {
|
||||
propBytes, err := json.Marshal(prop)
|
||||
if err == nil {
|
||||
var propInfo struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if json.Unmarshal(propBytes, &propInfo) == nil && propInfo.Type == "keyword" {
|
||||
e.useKeywordSuffix = false
|
||||
log.Infof("[Elasticsearch] Detected keyword type for ID fields, querying without .keyword suffix")
|
||||
return
|
||||
}
|
||||
}
|
||||
e.useKeywordSuffix = true
|
||||
log.Infof("[Elasticsearch] ID fields are not keyword type, querying with .keyword suffix")
|
||||
return
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] No mapping detected for chunk_id (empty index?), defaulting to .keyword suffix")
|
||||
e.useKeywordSuffix = true
|
||||
}
|
||||
|
||||
// EngineType returns the type of retriever engine (Elasticsearch)
|
||||
func (e *elasticsearchRepository) EngineType() typesLocal.RetrieverEngineType {
|
||||
return typesLocal.ElasticsearchRetrieverEngineType
|
||||
}
|
||||
|
||||
// Support returns the retrieval types supported by this repository (Keywords and Vector)
|
||||
func (e *elasticsearchRepository) Support() []typesLocal.RetrieverType {
|
||||
return []typesLocal.RetrieverType{typesLocal.KeywordsRetrieverType, typesLocal.VectorRetrieverType}
|
||||
}
|
||||
|
||||
// calculateStorageSize estimates the storage size in bytes for a single index document
|
||||
func (e *elasticsearchRepository) calculateStorageSize(embedding *elasticsearchRetriever.VectorEmbedding) int64 {
|
||||
// 1. Content text size
|
||||
contentSizeBytes := int64(len(embedding.Content))
|
||||
|
||||
// 2. Vector embedding size
|
||||
var vectorSizeBytes int64 = 0
|
||||
if embedding.Embedding != nil {
|
||||
// 4 bytes per dimension (full precision float)
|
||||
vectorSizeBytes = int64(len(embedding.Embedding) * 4)
|
||||
}
|
||||
|
||||
// 3. Metadata size (IDs, timestamps, and other fixed overhead)
|
||||
metadataSizeBytes := int64(250) // Approximately 250 bytes of metadata
|
||||
|
||||
// 4. Index overhead (Elasticsearch index expansion factor ~1.5)
|
||||
indexOverheadBytes := (contentSizeBytes + vectorSizeBytes) * 5 / 10
|
||||
|
||||
// Total size in bytes
|
||||
totalSizeBytes := contentSizeBytes + vectorSizeBytes + metadataSizeBytes + indexOverheadBytes
|
||||
return totalSizeBytes
|
||||
}
|
||||
|
||||
// EstimateStorageSize calculates the estimated storage size for a list of indices
|
||||
// Returns the total size in bytes
|
||||
func (e *elasticsearchRepository) EstimateStorageSize(ctx context.Context,
|
||||
indexInfoList []*typesLocal.IndexInfo, params map[string]any,
|
||||
) int64 {
|
||||
var totalStorageSize int64 = 0
|
||||
for _, embedding := range indexInfoList {
|
||||
embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, params)
|
||||
totalStorageSize += e.calculateStorageSize(embeddingDB)
|
||||
}
|
||||
logger.GetLogger(ctx).Infof(
|
||||
"[Elasticsearch] Storage size for %d indices: %d bytes", len(indexInfoList), totalStorageSize,
|
||||
)
|
||||
return totalStorageSize
|
||||
}
|
||||
|
||||
// Save stores a single index document in Elasticsearch
|
||||
// Returns an error if the operation fails
|
||||
func (e *elasticsearchRepository) Save(ctx context.Context,
|
||||
embedding *typesLocal.IndexInfo,
|
||||
additionalParams map[string]any,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
log.Debugf("[Elasticsearch] Saving index for chunk ID: %s", embedding.ChunkID)
|
||||
|
||||
// Convert to database format
|
||||
embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)
|
||||
if len(embeddingDB.Embedding) == 0 {
|
||||
err := fmt.Errorf("empty embedding vector for chunk ID: %s", embedding.ChunkID)
|
||||
log.Errorf("[Elasticsearch] %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// Index the document
|
||||
resp, err := e.client.Index(e.index).Request(embeddingDB).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to save index: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully saved index for chunk ID: %s, document ID: %s", embedding.ChunkID, resp.Id_)
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchSave stores multiple index documents in Elasticsearch using bulk API
|
||||
// Returns an error if the operation fails
|
||||
func (e *elasticsearchRepository) BatchSave(ctx context.Context,
|
||||
embeddingList []*typesLocal.IndexInfo, additionalParams map[string]any,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(embeddingList) == 0 {
|
||||
log.Warn("[Elasticsearch] Empty list provided to BatchSave, skipping")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Batch saving %d indices", len(embeddingList))
|
||||
indexRequest := e.client.Bulk().Index(e.index)
|
||||
|
||||
// Add each document to the bulk request
|
||||
for _, embedding := range embeddingList {
|
||||
embeddingDB := elasticsearchRetriever.ToDBVectorEmbedding(embedding, additionalParams)
|
||||
err := indexRequest.CreateOp(types.CreateOperation{Index_: &e.index}, embeddingDB)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to create bulk operation: %v", err)
|
||||
return fmt.Errorf("failed to create op: %w", err)
|
||||
}
|
||||
log.Debugf("[Elasticsearch] Added chunk ID %s to bulk request", embedding.ChunkID)
|
||||
}
|
||||
|
||||
// Execute the bulk request
|
||||
_, err := indexRequest.Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to execute bulk operation: %v", err)
|
||||
return fmt.Errorf("failed to do bulk: %w", err)
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully batch saved %d indices", len(embeddingList))
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteByChunkIDList removes documents from the index based on chunk IDs
|
||||
// Returns an error if the delete operation fails
|
||||
func (e *elasticsearchRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(chunkIDList) == 0 {
|
||||
log.Warn("[Elasticsearch] Empty chunk ID list provided for deletion, skipping")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Deleting indices by chunk IDs, count: %d", len(chunkIDList))
|
||||
// Use DeleteByQuery to delete all documents matching the chunk IDs
|
||||
_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
|
||||
Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("chunk_id"): chunkIDList}},
|
||||
}).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to delete by chunk IDs: %v", err)
|
||||
return fmt.Errorf("failed to delete by query: %w", err)
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully deleted documents by chunk IDs")
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteBySourceIDList removes documents from the index based on source IDs
|
||||
// Returns an error if the delete operation fails
|
||||
func (e *elasticsearchRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(sourceIDList) == 0 {
|
||||
log.Warn("[Elasticsearch] Empty source ID list provided for deletion, skipping")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Deleting indices by source IDs, count: %d", len(sourceIDList))
|
||||
// Use DeleteByQuery to delete all documents matching the source IDs
|
||||
_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
|
||||
Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("source_id"): sourceIDList}},
|
||||
}).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to delete by source IDs: %v", err)
|
||||
return fmt.Errorf("failed to delete by query: %w", err)
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully deleted documents by source IDs")
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteByKnowledgeIDList removes documents from the index based on knowledge IDs
|
||||
// Returns an error if the delete operation fails
|
||||
func (e *elasticsearchRepository) DeleteByKnowledgeIDList(ctx context.Context,
|
||||
knowledgeIDList []string, dimension int, knowledgeType string,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(knowledgeIDList) == 0 {
|
||||
log.Warn("[Elasticsearch] Empty knowledge ID list provided for deletion, skipping")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Deleting indices by knowledge IDs, count: %d", len(knowledgeIDList))
|
||||
// Use DeleteByQuery to delete all documents matching the knowledge IDs
|
||||
_, err := e.client.DeleteByQuery(e.index).Query(&types.Query{
|
||||
Terms: &types.TermsQuery{TermsQuery: map[string]types.TermsQueryField{e.idField("knowledge_id"): knowledgeIDList}},
|
||||
}).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to delete by knowledge IDs: %v", err)
|
||||
return fmt.Errorf("failed to delete by query: %w", err)
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully deleted documents by knowledge IDs")
|
||||
return nil
|
||||
}
|
||||
|
||||
// getBaseConds creates the base query conditions for retrieval operations
|
||||
// Returns a slice of Query objects with must and must_not conditions
|
||||
// KnowledgeBaseIDs and KnowledgeIDs use AND logic (search specific documents within knowledge bases)
|
||||
func (e *elasticsearchRepository) getBaseConds(params typesLocal.RetrieveParams) []types.Query {
|
||||
must := []types.Query{}
|
||||
|
||||
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
|
||||
// - If only KnowledgeBaseIDs: search entire knowledge bases
|
||||
// - If only KnowledgeIDs: search specific documents
|
||||
// - If both: search specific documents within the knowledge bases (AND)
|
||||
if len(params.KnowledgeBaseIDs) > 0 {
|
||||
must = append(must, types.Query{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{
|
||||
e.idField("knowledge_base_id"): params.KnowledgeBaseIDs,
|
||||
},
|
||||
}})
|
||||
}
|
||||
if len(params.KnowledgeIDs) > 0 {
|
||||
must = append(must, types.Query{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{
|
||||
e.idField("knowledge_id"): params.KnowledgeIDs,
|
||||
},
|
||||
}})
|
||||
}
|
||||
// Filter by tag IDs if specified
|
||||
if len(params.TagIDs) > 0 {
|
||||
must = append(must, types.Query{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{
|
||||
e.idField("tag_id"): params.TagIDs,
|
||||
},
|
||||
}})
|
||||
}
|
||||
|
||||
mustNot := make([]types.Query, 0)
|
||||
// Exclude disabled chunks (is_enabled = false)
|
||||
// Note: Historical data without is_enabled field will be included (not matching must_not)
|
||||
mustNot = append(mustNot, types.Query{Term: map[string]types.TermQuery{
|
||||
"is_enabled": {Value: false},
|
||||
}})
|
||||
if len(params.ExcludeKnowledgeIDs) > 0 {
|
||||
mustNot = append(mustNot, types.Query{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{e.idField("knowledge_id"): params.ExcludeKnowledgeIDs},
|
||||
}})
|
||||
}
|
||||
if len(params.ExcludeChunkIDs) > 0 {
|
||||
mustNot = append(mustNot, types.Query{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{e.idField("chunk_id"): params.ExcludeChunkIDs},
|
||||
}})
|
||||
}
|
||||
return []types.Query{{Bool: &types.BoolQuery{Must: must, MustNot: mustNot}}}
|
||||
}
|
||||
|
||||
// createIndexIfNotExists checks if the specified index exists and creates it if not
|
||||
// Returns an error if the operation fails
|
||||
func (e *elasticsearchRepository) createIndexIfNotExists(ctx context.Context) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
log.Debugf("[Elasticsearch] Checking if index exists: %s", e.index)
|
||||
|
||||
// Check if index exists
|
||||
exists, err := e.client.Indices.Exists(e.index).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to check if index exists: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
if exists {
|
||||
log.Debugf("[Elasticsearch] Index already exists: %s", e.index)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Create index if it doesn't exist, with optional shards/replicas settings
|
||||
log.Infof("[Elasticsearch] Creating index: %s", e.index)
|
||||
createReq := e.client.Indices.Create(e.index)
|
||||
if e.numberOfShards > 0 || e.numberOfReplicas >= 0 {
|
||||
settings := &types.IndexSettings{}
|
||||
if e.numberOfShards > 0 {
|
||||
shards := fmt.Sprintf("%d", e.numberOfShards)
|
||||
settings.NumberOfShards = &shards
|
||||
}
|
||||
if e.numberOfReplicas >= 0 {
|
||||
replicas := fmt.Sprintf("%d", e.numberOfReplicas)
|
||||
settings.NumberOfReplicas = &replicas
|
||||
}
|
||||
createReq = createReq.Settings(settings)
|
||||
}
|
||||
_, err = createReq.Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to create index: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Index created successfully: %s", e.index)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Retrieve dispatches the retrieval operation to the appropriate method based on retriever type
|
||||
// Returns a slice of RetrieveResult and an error if the operation fails
|
||||
func (e *elasticsearchRepository) Retrieve(ctx context.Context,
|
||||
params typesLocal.RetrieveParams,
|
||||
) ([]*typesLocal.RetrieveResult, error) {
|
||||
log := logger.GetLogger(ctx)
|
||||
log.Debugf("[Elasticsearch] Processing retrieval request of type: %s", params.RetrieverType)
|
||||
|
||||
// Route to appropriate retrieval method
|
||||
switch params.RetrieverType {
|
||||
case typesLocal.VectorRetrieverType:
|
||||
return e.VectorRetrieve(ctx, params)
|
||||
case typesLocal.KeywordsRetrieverType:
|
||||
return e.KeywordsRetrieve(ctx, params)
|
||||
}
|
||||
|
||||
err := fmt.Errorf("invalid retriever type: %v", params.RetrieverType)
|
||||
log.Errorf("[Elasticsearch] %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// VectorRetrieve performs vector similarity search using cosine similarity
|
||||
// Returns a slice of RetrieveResult containing matching documents
|
||||
func (e *elasticsearchRepository) VectorRetrieve(ctx context.Context,
|
||||
params typesLocal.RetrieveParams,
|
||||
) ([]*typesLocal.RetrieveResult, error) {
|
||||
log := logger.GetLogger(ctx)
|
||||
log.Infof("[Elasticsearch] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
|
||||
len(params.Embedding), params.TopK, params.Threshold)
|
||||
|
||||
filter := e.getBaseConds(params)
|
||||
|
||||
// Build script scoring query with cosine similarity
|
||||
queryVectorJSON, err := json.Marshal(params.Embedding)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to marshal query vector: %v", err)
|
||||
return nil, fmt.Errorf("failed to marshal query embedding: %w", err)
|
||||
}
|
||||
|
||||
scoreSource := "cosineSimilarity(params.query_vector, 'embedding')"
|
||||
minScore := float32(params.Threshold)
|
||||
scriptScore := &types.ScriptScoreQuery{
|
||||
Query: types.Query{Bool: &types.BoolQuery{Filter: filter}},
|
||||
Script: types.Script{
|
||||
Source: &scoreSource,
|
||||
Params: map[string]json.RawMessage{
|
||||
"query_vector": json.RawMessage(queryVectorJSON),
|
||||
},
|
||||
},
|
||||
MinScore: &minScore,
|
||||
}
|
||||
// Exclude embedding field from source to reduce response size
|
||||
sourceFilter := &types.SourceFilter{
|
||||
Excludes: []string{"embedding"},
|
||||
}
|
||||
|
||||
log.Debugf("[Elasticsearch] Executing vector search in index: %s", e.index)
|
||||
// Execute search with minimum score threshold
|
||||
response, err := e.client.Search().Index(e.index).Request(&search.Request{
|
||||
Query: &types.Query{ScriptScore: scriptScore},
|
||||
Size: ¶ms.TopK,
|
||||
Source_: sourceFilter,
|
||||
}).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Vector search failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Process search results
|
||||
var results []*typesLocal.IndexWithScore
|
||||
for _, hit := range response.Hits.Hits {
|
||||
var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
|
||||
if err := json.Unmarshal(hit.Source_, &embedding); err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to unmarshal search result: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
embedding.Score = float64(*hit.Score_)
|
||||
results = append(results,
|
||||
elasticsearchRetriever.FromDBVectorEmbeddingWithScore(*hit.Id_, embedding, typesLocal.MatchTypeEmbedding))
|
||||
}
|
||||
|
||||
if len(results) == 0 {
|
||||
log.Warnf("[Elasticsearch] No vector matches found that meet threshold %.4f", params.Threshold)
|
||||
} else {
|
||||
log.Infof("[Elasticsearch] Vector retrieval found %d results", len(results))
|
||||
log.Debugf("[Elasticsearch] Top result score: %.4f", results[0].Score)
|
||||
}
|
||||
|
||||
return []*typesLocal.RetrieveResult{
|
||||
{
|
||||
Results: results,
|
||||
RetrieverEngineType: typesLocal.ElasticsearchRetrieverEngineType,
|
||||
RetrieverType: typesLocal.VectorRetrieverType,
|
||||
Error: nil,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// KeywordsRetrieve performs keyword-based search in document content
|
||||
// Returns a slice of RetrieveResult containing matching documents
|
||||
func (e *elasticsearchRepository) KeywordsRetrieve(ctx context.Context,
|
||||
params typesLocal.RetrieveParams,
|
||||
) ([]*typesLocal.RetrieveResult, error) {
|
||||
log := logger.GetLogger(ctx)
|
||||
log.Infof("[Elasticsearch] Performing keywords retrieval with query: %s, topK: %d", params.Query, params.TopK)
|
||||
|
||||
filter := e.getBaseConds(params)
|
||||
// Build must conditions for content matching
|
||||
must := []types.Query{
|
||||
{Match: map[string]types.MatchQuery{"content": {Query: params.Query}}},
|
||||
}
|
||||
// Exclude embedding field from source to reduce response size
|
||||
sourceFilter := &types.SourceFilter{
|
||||
Excludes: []string{"embedding"},
|
||||
}
|
||||
|
||||
log.Debugf("[Elasticsearch] Executing keyword search in index: %s", e.index)
|
||||
response, err := e.client.Search().Index(e.index).Request(&search.Request{
|
||||
Query: &types.Query{Bool: &types.BoolQuery{Filter: filter, Must: must}},
|
||||
Size: ¶ms.TopK,
|
||||
Source_: sourceFilter,
|
||||
}).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Keywords search failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Process search results
|
||||
var results []*typesLocal.IndexWithScore
|
||||
for _, hit := range response.Hits.Hits {
|
||||
var embedding *elasticsearchRetriever.VectorEmbeddingWithScore
|
||||
if err := json.Unmarshal(hit.Source_, &embedding); err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to unmarshal search result: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
embedding.Score = float64(*hit.Score_)
|
||||
results = append(results,
|
||||
elasticsearchRetriever.FromDBVectorEmbeddingWithScore(*hit.Id_, embedding, typesLocal.MatchTypeKeywords),
|
||||
)
|
||||
}
|
||||
|
||||
if len(results) == 0 {
|
||||
log.Warnf("[Elasticsearch] No keyword matches found for query: %s", params.Query)
|
||||
} else {
|
||||
log.Infof("[Elasticsearch] Keywords retrieval found %d results", len(results))
|
||||
log.Debugf("[Elasticsearch] Top result score: %.4f", results[0].Score)
|
||||
}
|
||||
|
||||
return []*typesLocal.RetrieveResult{
|
||||
{
|
||||
Results: results,
|
||||
RetrieverEngineType: typesLocal.ElasticsearchRetrieverEngineType,
|
||||
RetrieverType: typesLocal.KeywordsRetrieverType,
|
||||
Error: nil,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// CopyIndices 复制索引数据
|
||||
func (e *elasticsearchRepository) CopyIndices(ctx context.Context,
|
||||
sourceKnowledgeBaseID string,
|
||||
sourceToTargetKBIDMap map[string]string,
|
||||
sourceToTargetChunkIDMap map[string]string,
|
||||
targetKnowledgeBaseID string,
|
||||
dimension int,
|
||||
knowledgeType string,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
log.Infof(
|
||||
"[Elasticsearch] Copying indices from source knowledge base %s to target knowledge base %s, count: %d",
|
||||
sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap),
|
||||
)
|
||||
|
||||
if len(sourceToTargetChunkIDMap) == 0 {
|
||||
log.Warn("[Elasticsearch] Empty mapping, skipping copy")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Build query parameters
|
||||
params := typesLocal.RetrieveParams{
|
||||
KnowledgeBaseIDs: []string{sourceKnowledgeBaseID},
|
||||
}
|
||||
|
||||
// Build base query conditions
|
||||
filter := e.getBaseConds(params)
|
||||
|
||||
// Set batch processing parameters
|
||||
batchSize := 500
|
||||
from := 0
|
||||
totalCopied := 0
|
||||
|
||||
for {
|
||||
// Execute pagination query
|
||||
searchResponse, err := e.client.Search().Index(e.index).
|
||||
Query(&types.Query{Bool: &types.BoolQuery{Filter: filter}}).
|
||||
From(from).
|
||||
Size(batchSize).
|
||||
Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to query source index data: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
hitsCount := len(searchResponse.Hits.Hits)
|
||||
if hitsCount == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Found %d source index data, batch start position: %d", hitsCount, from)
|
||||
|
||||
// Prepare index list for BatchSave
|
||||
var indexInfoList []*typesLocal.IndexInfo
|
||||
|
||||
// Collect all embedding vector data for additionalParams
|
||||
embeddingMap := make(map[string][]float32)
|
||||
|
||||
for _, hit := range searchResponse.Hits.Hits {
|
||||
// Parse source document
|
||||
var sourceDoc elasticsearchRetriever.VectorEmbedding
|
||||
if err := json.Unmarshal(hit.Source_, &sourceDoc); err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to parse source index data: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Get mapped target chunk ID and knowledge ID
|
||||
targetChunkID, ok := sourceToTargetChunkIDMap[sourceDoc.ChunkID]
|
||||
if !ok {
|
||||
log.Warnf("[Elasticsearch] Source chunk %s not found in target mapping, skipping", sourceDoc.ChunkID)
|
||||
continue
|
||||
}
|
||||
|
||||
targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceDoc.KnowledgeID]
|
||||
if !ok {
|
||||
log.Warnf(
|
||||
"[Elasticsearch] Source knowledge %s not found in target mapping, skipping",
|
||||
sourceDoc.KnowledgeID,
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// Save embedding vector to embeddingMap
|
||||
if len(sourceDoc.Embedding) > 0 {
|
||||
embeddingMap[targetChunkID] = sourceDoc.Embedding
|
||||
}
|
||||
|
||||
// Handle SourceID transformation for generated questions
|
||||
// Generated questions have SourceID format: {chunkID}-{questionID}
|
||||
// Regular chunks have SourceID == ChunkID
|
||||
var targetSourceID string
|
||||
if sourceDoc.SourceID == sourceDoc.ChunkID {
|
||||
// Regular chunk, use targetChunkID as SourceID
|
||||
targetSourceID = targetChunkID
|
||||
} else if strings.HasPrefix(sourceDoc.SourceID, sourceDoc.ChunkID+"-") {
|
||||
// This is a generated question, preserve the questionID part
|
||||
questionID := strings.TrimPrefix(sourceDoc.SourceID, sourceDoc.ChunkID+"-")
|
||||
targetSourceID = fmt.Sprintf("%s-%s", targetChunkID, questionID)
|
||||
} else {
|
||||
// For other complex scenarios, generate new unique SourceID
|
||||
targetSourceID = uuid.New().String()
|
||||
}
|
||||
|
||||
// Create new index information
|
||||
indexInfo := &typesLocal.IndexInfo{
|
||||
Content: sourceDoc.Content,
|
||||
SourceID: targetSourceID,
|
||||
SourceType: typesLocal.SourceType(sourceDoc.SourceType),
|
||||
ChunkID: targetChunkID,
|
||||
KnowledgeID: targetKnowledgeID,
|
||||
KnowledgeBaseID: targetKnowledgeBaseID,
|
||||
}
|
||||
|
||||
indexInfoList = append(indexInfoList, indexInfo)
|
||||
totalCopied++
|
||||
}
|
||||
|
||||
// Use BatchSave function to save index
|
||||
if len(indexInfoList) > 0 {
|
||||
// Add embedding vector to additional parameters
|
||||
additionalParams := map[string]any{
|
||||
"embedding": embeddingMap,
|
||||
}
|
||||
|
||||
if err := e.BatchSave(ctx, indexInfoList, additionalParams); err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to batch save index: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully copied batch data, batch size: %d, total copied: %d",
|
||||
len(indexInfoList), totalCopied)
|
||||
}
|
||||
|
||||
// Move to next batch
|
||||
from += hitsCount
|
||||
|
||||
// If the number of returned records is less than the request size, it means the last page has been reached
|
||||
if hitsCount < batchSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Index copy completed, total copied: %d", totalCopied)
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
|
||||
func (e *elasticsearchRepository) BatchUpdateChunkEnabledStatus(
|
||||
ctx context.Context,
|
||||
chunkStatusMap map[string]bool,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(chunkStatusMap) == 0 {
|
||||
log.Warnf("[Elasticsearch] Chunk status map is empty, skipping update")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Batch updating chunk enabled status, count: %d", len(chunkStatusMap))
|
||||
|
||||
// Group chunks by enabled status for batch updates
|
||||
enabledChunkIDs := make([]string, 0)
|
||||
disabledChunkIDs := make([]string, 0)
|
||||
|
||||
for chunkID, enabled := range chunkStatusMap {
|
||||
if enabled {
|
||||
enabledChunkIDs = append(enabledChunkIDs, chunkID)
|
||||
} else {
|
||||
disabledChunkIDs = append(disabledChunkIDs, chunkID)
|
||||
}
|
||||
}
|
||||
|
||||
// Batch update enabled chunks using update_by_query
|
||||
if len(enabledChunkIDs) > 0 {
|
||||
query := types.NewQuery()
|
||||
query.Bool = &types.BoolQuery{
|
||||
Must: []types.Query{
|
||||
{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{
|
||||
e.idField("chunk_id"): enabledChunkIDs,
|
||||
},
|
||||
}},
|
||||
},
|
||||
}
|
||||
source := "ctx._source.is_enabled = true"
|
||||
lang := scriptlanguage.Painless
|
||||
script := types.Script{
|
||||
Source: &source,
|
||||
Lang: &lang,
|
||||
}
|
||||
_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to update enabled chunks: %v", err)
|
||||
return err
|
||||
}
|
||||
log.Infof("[Elasticsearch] Updated %d chunks to enabled", len(enabledChunkIDs))
|
||||
}
|
||||
|
||||
// Batch update disabled chunks using update_by_query
|
||||
if len(disabledChunkIDs) > 0 {
|
||||
query := types.NewQuery()
|
||||
query.Bool = &types.BoolQuery{
|
||||
Must: []types.Query{
|
||||
{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{
|
||||
e.idField("chunk_id"): disabledChunkIDs,
|
||||
},
|
||||
}},
|
||||
},
|
||||
}
|
||||
source := "ctx._source.is_enabled = false"
|
||||
lang := scriptlanguage.Painless
|
||||
script := types.Script{
|
||||
Source: &source,
|
||||
Lang: &lang,
|
||||
}
|
||||
_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to update disabled chunks: %v", err)
|
||||
return err
|
||||
}
|
||||
log.Infof("[Elasticsearch] Updated %d chunks to disabled", len(disabledChunkIDs))
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully batch updated chunk enabled status")
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchUpdateChunkTagID updates the tag ID of chunks in batch
|
||||
func (e *elasticsearchRepository) BatchUpdateChunkTagID(
|
||||
ctx context.Context,
|
||||
chunkTagMap map[string]string,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(chunkTagMap) == 0 {
|
||||
log.Warnf("[Elasticsearch] Chunk tag map is empty, skipping update")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Batch updating chunk tag ID, count: %d", len(chunkTagMap))
|
||||
|
||||
// Group chunks by tag ID for batch updates
|
||||
tagGroups := make(map[string][]string)
|
||||
for chunkID, tagID := range chunkTagMap {
|
||||
tagGroups[tagID] = append(tagGroups[tagID], chunkID)
|
||||
}
|
||||
|
||||
// Batch update chunks for each tag ID using update_by_query
|
||||
for tagID, chunkIDs := range tagGroups {
|
||||
query := types.NewQuery()
|
||||
query.Bool = &types.BoolQuery{
|
||||
Must: []types.Query{
|
||||
{Terms: &types.TermsQuery{
|
||||
TermsQuery: map[string]types.TermsQueryField{
|
||||
e.idField("chunk_id"): chunkIDs,
|
||||
},
|
||||
}},
|
||||
},
|
||||
}
|
||||
source := "ctx._source.tag_id = params.tag_id"
|
||||
lang := scriptlanguage.Painless
|
||||
script := types.Script{
|
||||
Source: &source,
|
||||
Lang: &lang,
|
||||
Params: map[string]json.RawMessage{
|
||||
"tag_id": json.RawMessage(`"` + tagID + `"`),
|
||||
},
|
||||
}
|
||||
_, err := e.client.UpdateByQuery(e.index).Query(query).Script(&script).Do(ctx)
|
||||
if err != nil {
|
||||
log.Errorf("[Elasticsearch] Failed to update chunks with tag_id %s: %v", tagID, err)
|
||||
return err
|
||||
}
|
||||
log.Infof("[Elasticsearch] Updated %d chunks to tag_id=%s", len(chunkIDs), tagID)
|
||||
}
|
||||
|
||||
log.Infof("[Elasticsearch] Successfully batch updated chunk tag ID")
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,455 @@
|
||||
package legacygraph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
neo4jrepo "github.com/Tencent/WeKnora/internal/application/repository/retriever/neo4j"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/neo4j/neo4j-go-driver/v6/neo4j"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Repository preserves source records while exposing verified native chunks to RAG.
|
||||
type Repository struct {
|
||||
db *gorm.DB
|
||||
fallback interfaces.RetrieveGraphRepository
|
||||
}
|
||||
|
||||
func NewRepository(db *gorm.DB, driver neo4j.Driver) *Repository {
|
||||
r := &Repository{db: db}
|
||||
if driver != nil {
|
||||
r.fallback = neo4jrepo.NewNeo4jRepository(driver)
|
||||
}
|
||||
return r
|
||||
}
|
||||
func NewRetrieveRepository(db *gorm.DB, driver neo4j.Driver) interfaces.RetrieveGraphRepository {
|
||||
return NewRepository(db, driver)
|
||||
}
|
||||
func (r *Repository) AddGraph(ctx context.Context, ns types.NameSpace, graphs []*types.GraphData) error {
|
||||
return r.addNativeGraphs(ctx, ns, graphs, nil)
|
||||
}
|
||||
func (r *Repository) DelGraph(ctx context.Context, ns []types.NameSpace) error {
|
||||
return r.deleteNativeGraphs(ctx, ns)
|
||||
}
|
||||
|
||||
type Record struct {
|
||||
ID int64 `json:"id" gorm:"column:legacy_id"`
|
||||
Payload json.RawMessage `json:"payload" gorm:"column:payload"`
|
||||
}
|
||||
type Evidence struct {
|
||||
Record
|
||||
NativeKnowledgeID string `json:"native_knowledge_id"`
|
||||
NativeChunkID string `json:"native_chunk_id"`
|
||||
MappingState string `json:"mapping_state"`
|
||||
}
|
||||
type Overview struct {
|
||||
Counts map[string]int64 `json:"counts"`
|
||||
Sources int64 `json:"sources"`
|
||||
MappedSources int64 `json:"mapped_sources"`
|
||||
Nodes []Record `json:"nodes"`
|
||||
Relations []Record `json:"relations"`
|
||||
Evidence []Evidence `json:"evidence"`
|
||||
NextCursor int64 `json:"next_cursor"`
|
||||
}
|
||||
|
||||
func (r *Repository) scope(ctx context.Context, kb string) (*gorm.DB, error) {
|
||||
tenant, ok := types.TenantIDFromContext(ctx)
|
||||
if !ok || tenant == 0 || kb == "" {
|
||||
return nil, fmt.Errorf("graph tenant and knowledge base are required")
|
||||
}
|
||||
return r.db.WithContext(ctx).Table("native_graph_records").Where("tenant_id = ? AND knowledge_base_id = ?", tenant, kb), nil
|
||||
}
|
||||
func (r *Repository) Browse(ctx context.Context, kb, query string, cursor, entity int64) (*Overview, error) {
|
||||
scope, err := r.scope(ctx, kb)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := &Overview{Counts: map[string]int64{}, Nodes: []Record{}, Relations: []Record{}, Evidence: []Evidence{}}
|
||||
var counts []struct {
|
||||
SourceTable string
|
||||
N int64
|
||||
}
|
||||
if err = scope.Session(&gorm.Session{}).Select("source_table, count(*) AS n").Group("source_table").Scan(&counts).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, c := range counts {
|
||||
out.Counts[c.SourceTable] = c.N
|
||||
}
|
||||
tenant := types.MustTenantIDFromContext(ctx)
|
||||
var coverage struct {
|
||||
Sources int64
|
||||
Mapped int64
|
||||
}
|
||||
// Overview coverage records migration-time matches; each returned link is revalidated below.
|
||||
err = r.db.WithContext(ctx).Raw(`SELECT count(*) AS sources, count(native_chunk_id) AS mapped FROM native_graph_sources WHERE tenant_id=? AND knowledge_base_id=?`, tenant, kb).Scan(&coverage).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out.Sources = coverage.Sources
|
||||
out.MappedSources = coverage.Mapped
|
||||
q := scope.Session(&gorm.Session{}).Select("legacy_id, payload").Where("source_table = ? AND legacy_id > ?", "knowledgeentity", cursor)
|
||||
if entity > 0 {
|
||||
q = q.Where("legacy_id = ?", entity)
|
||||
} else if strings.TrimSpace(query) != "" {
|
||||
q = q.Where("strpos(lower(payload->>'canonical_name'), lower(?)) > 0 OR strpos(lower(?), lower(payload->>'canonical_name')) > 0", query, query)
|
||||
}
|
||||
if err = q.Order("legacy_id").Limit(51).Find(&out.Nodes).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(out.Nodes) > 50 {
|
||||
out.NextCursor = out.Nodes[49].ID
|
||||
out.Nodes = out.Nodes[:50]
|
||||
}
|
||||
if err = r.projectNativeEntityAttributes(ctx, types.NameSpace{KnowledgeBase: kb}, out.Nodes, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if entity <= 0 {
|
||||
return out, nil
|
||||
}
|
||||
err = r.db.WithContext(ctx).Raw(`SELECT r.legacy_id, r.payload || jsonb_build_object('from_name',a.payload->>'canonical_name','to_name',b.payload->>'canonical_name') AS payload FROM native_graph_records r LEFT JOIN native_graph_records a ON a.tenant_id=r.tenant_id AND a.knowledge_base_id=r.knowledge_base_id AND a.source_table='knowledgeentity' AND a.legacy_id=(r.payload->>'from_entity_id')::bigint LEFT JOIN native_graph_records b ON b.tenant_id=r.tenant_id AND b.knowledge_base_id=r.knowledge_base_id AND b.source_table='knowledgeentity' AND b.legacy_id=(r.payload->>'to_entity_id')::bigint WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgerelation' AND (r.payload->>'from_entity_id'=? OR r.payload->>'to_entity_id'=?) ORDER BY r.legacy_id LIMIT 100`, tenant, kb, strconv.FormatInt(entity, 10), strconv.FormatInt(entity, 10)).Scan(&out.Relations).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
err = r.db.WithContext(ctx).Raw(`SELECT r.legacy_id,r.payload, COALESCE(c.knowledge_id,'') AS native_knowledge_id, COALESCE(c.id,'') AS native_chunk_id, CASE WHEN c.id IS NOT NULL THEN s.mapping_state WHEN s.native_chunk_id IS NOT NULL THEN 'stale' ELSE COALESCE(s.mapping_state,'unmapped') END AS mapping_state FROM native_graph_records r LEFT JOIN native_graph_sources s ON s.tenant_id=r.tenant_id AND s.knowledge_base_id=r.knowledge_base_id AND s.source_chunk_id=r.payload->>'source_chunk_id' LEFT JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256 AND EXISTS (SELECT 1 FROM knowledges k WHERE k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL) AND EXISTS (SELECT 1 FROM knowledge_bases kb WHERE kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL) WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgegraphevidence' AND r.payload->>'subject_type'='entity' AND r.payload->>'subject_id'=? ORDER BY r.legacy_id LIMIT 100`, tenant, kb, strconv.FormatInt(entity, 10)).Scan(&out.Evidence).Error
|
||||
return out, err
|
||||
}
|
||||
|
||||
func (r *Repository) SearchNode(ctx context.Context, ns types.NameSpace, terms []string) (*types.GraphData, error) {
|
||||
out := &types.GraphData{Node: []*types.GraphNode{}, Relation: []*types.GraphRelation{}}
|
||||
scope, err := r.scope(ctx, ns.KnowledgeBase)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
terms = boundedTerms(terms)
|
||||
if len(terms) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
clauses := make([]string, 0, len(terms))
|
||||
args := make([]any, 0, len(terms)*2)
|
||||
for _, term := range terms {
|
||||
clauses = append(clauses, "(strpos(lower(payload->>'canonical_name'),lower(?))>0 OR strpos(lower(?),lower(payload->>'canonical_name'))>0)")
|
||||
args = append(args, term, term)
|
||||
}
|
||||
var nodes []Record
|
||||
err = scope.Session(&gorm.Session{}).Select("legacy_id,payload").Where("source_table='knowledgeentity' AND payload->>'status'='active'").Where(strings.Join(clauses, " OR "), args...).Order("legacy_id").Limit(16).Find(&nodes).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tenant := types.MustTenantIDFromContext(ctx)
|
||||
seedIDs := []string{}
|
||||
seenIDs := map[int64]bool{}
|
||||
for _, node := range nodes {
|
||||
seedIDs = append(seedIDs, strconv.FormatInt(node.ID, 10))
|
||||
seenIDs[node.ID] = true
|
||||
}
|
||||
var relations []Record
|
||||
if len(seedIDs) > 0 {
|
||||
q, _ := r.scope(ctx, ns.KnowledgeBase)
|
||||
err = q.Select("legacy_id,payload").Where("source_table='knowledgerelation' AND payload->>'status'='active' AND (payload->>'from_entity_id' IN ? OR payload->>'to_entity_id' IN ?)", seedIDs, seedIDs).Order("legacy_id").Limit(32).Find(&relations).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
neighborIDs := []int64{}
|
||||
for _, rel := range relations {
|
||||
var edge struct {
|
||||
From int64 `json:"from_entity_id"`
|
||||
To int64 `json:"to_entity_id"`
|
||||
}
|
||||
if err = json.Unmarshal(rel.Payload, &edge); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, id := range []int64{edge.From, edge.To} {
|
||||
if !seenIDs[id] && len(neighborIDs) < 16 {
|
||||
seenIDs[id] = true
|
||||
neighborIDs = append(neighborIDs, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(neighborIDs) > 0 {
|
||||
var neighbors []Record
|
||||
q, _ = r.scope(ctx, ns.KnowledgeBase)
|
||||
if err = q.Select("legacy_id,payload").Where("source_table='knowledgeentity' AND payload->>'status'='active' AND legacy_id IN ?", neighborIDs).Order("legacy_id").Find(&neighbors).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
nodes = append(nodes, neighbors...)
|
||||
}
|
||||
}
|
||||
chunksByEntity := map[int64][]string{}
|
||||
chunksByRelation := map[int64][]string{}
|
||||
if len(nodes) > 0 {
|
||||
type subject struct {
|
||||
ID int64 `json:"subject_id"`
|
||||
Type string `json:"subject_type"`
|
||||
}
|
||||
ids := make([]subject, 0, len(nodes)+len(relations))
|
||||
for _, node := range nodes {
|
||||
ids = append(ids, subject{node.ID, "entity"})
|
||||
}
|
||||
for _, relation := range relations {
|
||||
ids = append(ids, subject{relation.ID, "relation"})
|
||||
}
|
||||
encodedIDs, err := json.Marshal(ids)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var matches []struct {
|
||||
SubjectID int64
|
||||
SubjectType string
|
||||
ChunkID string
|
||||
}
|
||||
// One bounded database request, with the evidence index and four-chunk cap per entity.
|
||||
query := `SELECT seed.subject_id,seed.subject_type,matched.chunk_id FROM jsonb_to_recordset(?::jsonb) AS seed(subject_id bigint,subject_type text) CROSS JOIN LATERAL (
|
||||
SELECT DISTINCT c.id AS chunk_id FROM native_graph_records e
|
||||
JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
|
||||
JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
|
||||
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
|
||||
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
|
||||
WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'=seed.subject_type AND e.payload->>'subject_id'=seed.subject_id::text AND e.payload->>'status'='active'`
|
||||
args := []any{string(encodedIDs), tenant, ns.KnowledgeBase}
|
||||
if ns.Knowledge != "" {
|
||||
query += " AND c.knowledge_id=?"
|
||||
args = append(args, ns.Knowledge)
|
||||
}
|
||||
query += " ORDER BY c.id LIMIT 4) matched"
|
||||
if err = r.db.WithContext(ctx).Raw(query, args...).Scan(&matches).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, match := range matches {
|
||||
if match.SubjectType == "entity" {
|
||||
chunksByEntity[match.SubjectID] = append(chunksByEntity[match.SubjectID], match.ChunkID)
|
||||
} else {
|
||||
chunksByRelation[match.SubjectID] = append(chunksByRelation[match.SubjectID], match.ChunkID)
|
||||
}
|
||||
}
|
||||
}
|
||||
if err = r.projectNativeEntityAttributes(ctx, ns, nodes, chunksByEntity); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
names := map[int64]string{}
|
||||
for _, node := range nodes {
|
||||
var entity struct {
|
||||
Name string `json:"canonical_name"`
|
||||
Type string `json:"entity_type"`
|
||||
Attributes []string `json:"attributes"`
|
||||
}
|
||||
if err = json.Unmarshal(node.Payload, &entity); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
chunks := chunksByEntity[node.ID]
|
||||
if len(chunks) > 0 {
|
||||
names[node.ID] = entity.Name
|
||||
attributes := []string{entity.Type}
|
||||
identity := fmt.Sprintf("legacy_entity:%d", node.ID)
|
||||
for _, attribute := range entity.Attributes {
|
||||
if attribute != entity.Type && attribute != identity {
|
||||
attributes = append(attributes, attribute)
|
||||
}
|
||||
}
|
||||
attributes = append(attributes, identity)
|
||||
out.Node = append(out.Node, &types.GraphNode{Name: entity.Name, Chunks: chunks, Attributes: attributes})
|
||||
}
|
||||
}
|
||||
for _, rel := range relations {
|
||||
var edge struct {
|
||||
From int64 `json:"from_entity_id"`
|
||||
To int64 `json:"to_entity_id"`
|
||||
Type string `json:"relation_type"`
|
||||
Origin string `json:"origin"`
|
||||
}
|
||||
if err = json.Unmarshal(rel.Payload, &edge); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if edge.Origin == nativeOrigin && len(chunksByRelation[rel.ID]) == 0 {
|
||||
continue
|
||||
}
|
||||
if names[edge.From] != "" && names[edge.To] != "" {
|
||||
out.Relation = append(out.Relation, &types.GraphRelation{Node1: names[edge.From], Node2: names[edge.To], Type: edge.Type})
|
||||
}
|
||||
}
|
||||
if r.fallback != nil {
|
||||
fallback, err := r.fallback.SearchNode(ctx, ns, terms)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if fallback != nil {
|
||||
out.Node = append(out.Node, fallback.Node...)
|
||||
out.Relation = append(out.Relation, fallback.Relation...)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Attribute ownership is checked independently of the four returned chunks:
|
||||
// an entity may still have a bare reference after its attribute source expires.
|
||||
func (r *Repository) projectNativeEntityAttributes(ctx context.Context, ns types.NameSpace, nodes []Record, chunks map[int64][]string) error {
|
||||
type subject struct {
|
||||
ID int64 `json:"subject_id"`
|
||||
Knowledge string `json:"knowledge_id"`
|
||||
Type string `json:"entity_type"`
|
||||
Attributes json.RawMessage `json:"attributes"`
|
||||
}
|
||||
seeds := []subject{}
|
||||
payloads := map[int64]map[string]json.RawMessage{}
|
||||
for _, node := range nodes {
|
||||
var entity struct {
|
||||
Origin string `json:"origin"`
|
||||
Knowledge string `json:"native_knowledge_id"`
|
||||
Type string `json:"entity_type"`
|
||||
Attributes json.RawMessage `json:"attributes"`
|
||||
}
|
||||
if err := json.Unmarshal(node.Payload, &entity); err != nil {
|
||||
return err
|
||||
}
|
||||
if entity.Origin != nativeOrigin {
|
||||
continue
|
||||
}
|
||||
var payload map[string]json.RawMessage
|
||||
if err := json.Unmarshal(node.Payload, &payload); err != nil {
|
||||
return err
|
||||
}
|
||||
payloads[node.ID] = payload
|
||||
seeds = append(seeds, subject{node.ID, entity.Knowledge, entity.Type, entity.Attributes})
|
||||
}
|
||||
if len(seeds) == 0 {
|
||||
return nil
|
||||
}
|
||||
encoded, err := json.Marshal(seeds)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var matches []struct {
|
||||
SubjectID int64
|
||||
Snapshots json.RawMessage
|
||||
Chunks json.RawMessage
|
||||
AttributeChunk string
|
||||
}
|
||||
query := `SELECT seed.subject_id,jsonb_agg(DISTINCT e.payload->'entity_attributes') AS snapshots,
|
||||
to_jsonb((array_agg(DISTINCT c.id ORDER BY c.id))[1:4]) AS chunks,
|
||||
COALESCE(min(c.id) FILTER (WHERE e.payload->'entity_attributes'->>0 IS NOT NULL
|
||||
AND (e.payload->'entity_attributes'=seed.attributes OR e.payload->'entity_attributes'->>0=seed.entity_type)),'') AS attribute_chunk
|
||||
FROM jsonb_to_recordset(?::jsonb) AS seed(subject_id bigint,knowledge_id text,entity_type text,attributes jsonb)
|
||||
JOIN native_graph_records e ON e.payload->>'subject_id'=seed.subject_id::text AND e.payload->>'native_knowledge_id'=seed.knowledge_id
|
||||
JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
|
||||
JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id
|
||||
AND s.native_knowledge_id=c.knowledge_id AND c.deleted_at IS NULL AND c.is_enabled=true
|
||||
AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
|
||||
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
|
||||
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
|
||||
WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'='entity'
|
||||
AND e.payload->>'origin'=? AND e.payload->>'status'='active' AND s.mapping_state='native_extraction'
|
||||
AND e.payload->>'source_content_hash'=s.native_content_sha256 AND e.source_snapshot_sha256=s.native_content_sha256
|
||||
AND s.source_content_hash=s.native_content_sha256`
|
||||
args := []any{string(encoded), types.MustTenantIDFromContext(ctx), ns.KnowledgeBase, nativeOrigin}
|
||||
if ns.Knowledge != "" {
|
||||
query += " AND c.knowledge_id=?"
|
||||
args = append(args, ns.Knowledge)
|
||||
}
|
||||
query += " GROUP BY seed.subject_id"
|
||||
if err = r.db.WithContext(ctx).Raw(query, args...).Scan(&matches).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
supported := map[int64][]string{}
|
||||
attributeChunks := map[int64]string{}
|
||||
verifiedChunks := map[int64][]string{}
|
||||
for _, match := range matches {
|
||||
supported[match.SubjectID] = compatibleNativeAttributeSnapshots(match.Snapshots)
|
||||
attributeChunks[match.SubjectID] = match.AttributeChunk
|
||||
var current []string
|
||||
if err = json.Unmarshal(match.Chunks, ¤t); err != nil {
|
||||
return err
|
||||
}
|
||||
verifiedChunks[match.SubjectID] = current
|
||||
}
|
||||
for i, node := range nodes {
|
||||
payload, native := payloads[node.ID]
|
||||
if !native {
|
||||
continue
|
||||
}
|
||||
if chunks != nil {
|
||||
chunks[node.ID] = verifiedChunks[node.ID]
|
||||
}
|
||||
attributes := []string{}
|
||||
entityType := "concept"
|
||||
if snapshot := supported[node.ID]; len(snapshot) > 0 {
|
||||
var stored []string
|
||||
if err := json.Unmarshal(payload["attributes"], &stored); err == nil && slices.Equal(stored, snapshot) {
|
||||
attributes = stored
|
||||
entityType = snapshot[0]
|
||||
}
|
||||
var storedType string
|
||||
if err := json.Unmarshal(payload["entity_type"], &storedType); err == nil && storedType == snapshot[0] {
|
||||
entityType = storedType
|
||||
}
|
||||
}
|
||||
if chunk := attributeChunks[node.ID]; chunks != nil && chunk != "" && (len(attributes) > 0 || entityType != "concept") {
|
||||
verified := []string{chunk}
|
||||
for _, prior := range chunks[node.ID] {
|
||||
if prior != chunk && len(verified) < 4 {
|
||||
verified = append(verified, prior)
|
||||
}
|
||||
}
|
||||
chunks[node.ID] = verified
|
||||
}
|
||||
payload["attributes"], _ = json.Marshal(attributes)
|
||||
payload["entity_type"], _ = json.Marshal(entityType)
|
||||
if nodes[i].Payload, err = json.Marshal(payload); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func compatibleNativeAttributeSnapshots(raw json.RawMessage) []string {
|
||||
var snapshots []json.RawMessage
|
||||
if err := json.Unmarshal(raw, &snapshots); err != nil {
|
||||
return nil
|
||||
}
|
||||
var attributes []string
|
||||
for _, snapshot := range snapshots {
|
||||
var values []*string
|
||||
if err := json.Unmarshal(snapshot, &values); err != nil {
|
||||
return nil
|
||||
}
|
||||
if len(values) == 0 {
|
||||
continue
|
||||
}
|
||||
current := make([]string, len(values))
|
||||
for i, value := range values {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
current[i] = *value
|
||||
}
|
||||
if len(attributes) > 0 && !slices.Equal(attributes, current) {
|
||||
return nil
|
||||
}
|
||||
attributes = current
|
||||
}
|
||||
return attributes
|
||||
}
|
||||
|
||||
func boundedTerms(terms []string) []string {
|
||||
out := []string{}
|
||||
seen := map[string]bool{}
|
||||
for _, term := range terms {
|
||||
term = strings.TrimSpace(term)
|
||||
if term == "" || seen[term] {
|
||||
continue
|
||||
}
|
||||
if len([]rune(term)) > 200 {
|
||||
term = string([]rune(term)[:200])
|
||||
}
|
||||
seen[term] = true
|
||||
out = append(out, term)
|
||||
if len(out) == 8 {
|
||||
break
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package legacygraph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestSearchNodePreservesStoredAttributes(t *testing.T) {
|
||||
for _, origin := range []string{"native", "legacy"} {
|
||||
for _, stored := range []bool{true, false} {
|
||||
t.Run(fmt.Sprintf("%s/attributes=%t", origin, stored), func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
graph := testGraph("chunk")
|
||||
graph.Node[0].Attributes = []string{"concept", "Not a medicine", "No treatment, nursing or health-care effect", "Evidence keeps spaces\nand line breaks", "concept"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if origin == "legacy" {
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'origin' WHERE source_table='knowledgeentity'`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Exec(`UPDATE native_graph_sources SET mapping_state='exact_content_hash'`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if !stored {
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'attributes' WHERE source_table='knowledgeentity'`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
want := []string{"concept"}
|
||||
if stored {
|
||||
want = append(want, graph.Node[0].Attributes[1:4]...)
|
||||
}
|
||||
want = append(want, fmt.Sprintf("legacy_entity:%d", id))
|
||||
// The same record must preserve attributes as either a seed or a neighbor.
|
||||
for _, term := range []string{"Trust", "Learning"} {
|
||||
found, err := r.SearchNode(ctx, ns, []string{term})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got *types.GraphNode
|
||||
for _, node := range found.Node {
|
||||
if node.Name == "Trust" {
|
||||
got = node
|
||||
}
|
||||
}
|
||||
if got == nil || !reflect.DeepEqual(got.Attributes, want) || !reflect.DeepEqual(got.Chunks, []string{"chunk"}) {
|
||||
t.Fatalf("%s lost source attributes or compatibility metadata: got=%#v want=%#v", term, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchNodeAttributesRespectSourceAvailabilityAndOwnership(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
graph := testGraph("chunk")
|
||||
graph.Node[0].Attributes = []string{"concept", "Not a medicine", "No treatment, nursing or health-care effect"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, scope := range []struct {
|
||||
name string
|
||||
tenant uint64
|
||||
ns types.NameSpace
|
||||
}{
|
||||
{"other tenant", 8, ns},
|
||||
{"other knowledge base", 7, types.NameSpace{KnowledgeBase: "other-kb"}},
|
||||
{"other document", 7, types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc2"}},
|
||||
} {
|
||||
t.Run(scope.name, func(t *testing.T) {
|
||||
otherCtx := context.WithValue(ctx, types.TenantIDContextKey, scope.tenant)
|
||||
got, err := r.SearchNode(otherCtx, scope.ns, []string{"Trust"})
|
||||
if err != nil || len(got.Node) != 0 || len(got.Relation) != 0 {
|
||||
t.Fatalf("out-of-scope attributes returned: %#v %v", got, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, change := range []struct{ name, sql string }{
|
||||
{"edited source", `UPDATE chunks SET content=content || ' changed' WHERE id='chunk'`},
|
||||
{"disabled source", `UPDATE chunks SET is_enabled=false WHERE id='chunk'`},
|
||||
{"deleted source", `UPDATE chunks SET deleted_at=now() WHERE id='chunk'`},
|
||||
{"deleted document", `UPDATE knowledges SET deleted_at=now() WHERE id='doc' AND tenant_id=7`},
|
||||
{"deleted knowledge base", `UPDATE knowledge_bases SET deleted_at=now() WHERE id='kb' AND tenant_id=7`},
|
||||
} {
|
||||
t.Run(change.name, func(t *testing.T) {
|
||||
if err := tx.SavePoint("attribute_source_check").Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer tx.RollbackTo("attribute_source_check")
|
||||
if err := tx.Exec(change.sql).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := r.SearchNode(ctx, ns, []string{"Trust"})
|
||||
if err != nil || len(got.Node) != 0 || len(got.Relation) != 0 {
|
||||
t.Fatalf("unavailable source attributes returned: %#v %v", got, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
package legacygraph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestMissingGraphScope(t *testing.T) {
|
||||
r := NewRepository(nil, nil)
|
||||
if _, err := r.Browse(context.Background(), "kb", "", 0, 0); err == nil {
|
||||
t.Fatal("missing tenant accepted")
|
||||
}
|
||||
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(10059))
|
||||
if _, err := r.SearchNode(ctx, types.NameSpace{}, []string{"term"}); err == nil {
|
||||
t.Fatal("missing KB accepted")
|
||||
}
|
||||
}
|
||||
func TestBoundedTerms(t *testing.T) {
|
||||
terms := boundedTerms([]string{"", " ", "a", "a", strings.Repeat("中", 240), "b", "c", "d", "e", "f", "g", "h"})
|
||||
if len(terms) != 8 || terms[0] != "a" || len([]rune(terms[1])) != 200 {
|
||||
t.Fatalf("unbounded terms: %#v", terms)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchNodeKeepsGenericSourcesIsolated(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Evidence about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
for chunk, target := range map[string]string{"chunk": "First target", "sibling": "Second target"} {
|
||||
graph := &types.GraphData{
|
||||
Node: []*types.GraphNode{
|
||||
{Name: "本品", Chunks: []string{chunk}},
|
||||
{Name: target, Chunks: []string{chunk}},
|
||||
},
|
||||
Relation: []*types.GraphRelation{{Node1: "本品", Node2: target, Type: "关联"}},
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
for chunk, target := range map[string]string{"chunk": "First target", "sibling": "Second target"} {
|
||||
graph, err := r.SearchNode(ctx, ns, []string{target})
|
||||
if err != nil || len(graph.Node) != 2 || len(graph.Relation) != 1 {
|
||||
t.Fatalf("unexpected single-product result: %#v %v", graph, err)
|
||||
}
|
||||
for _, node := range graph.Node {
|
||||
if len(node.Chunks) != 1 || node.Chunks[0] != chunk {
|
||||
t.Fatalf("same-name neighbor expanded another product: %#v", node)
|
||||
}
|
||||
}
|
||||
if edge := graph.Relation[0]; edge.Node1 != "本品" || edge.Node2 != target {
|
||||
t.Fatalf("unexpected relation projection: %#v", edge)
|
||||
}
|
||||
}
|
||||
graph, err := r.SearchNode(ctx, ns, []string{"本品"})
|
||||
if err != nil || len(graph.Node) != 4 || len(graph.Relation) != 2 {
|
||||
t.Fatalf("same-name nodes were collapsed: %#v %v", graph, err)
|
||||
}
|
||||
sources := map[string]bool{}
|
||||
identities := map[string]bool{}
|
||||
for _, node := range graph.Node {
|
||||
if node.Name != "本品" {
|
||||
continue
|
||||
}
|
||||
if len(node.Chunks) != 1 || len(node.Attributes) != 2 {
|
||||
t.Fatalf("generic source identity was lost: %#v", node)
|
||||
}
|
||||
sources[node.Chunks[0]] = true
|
||||
identities[node.Attributes[1]] = true
|
||||
}
|
||||
if len(sources) != 2 || !sources["chunk"] || !sources["sibling"] || len(identities) != 2 {
|
||||
t.Fatalf("generic products share retrieval identity: sources=%v identities=%v", sources, identities)
|
||||
}
|
||||
}
|
||||
|
||||
func assertReaderTrustAttributes(t *testing.T, r *Repository, ctx context.Context, want []string) {
|
||||
t.Helper()
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
wantType := "concept"
|
||||
if len(want) > 0 {
|
||||
wantType = want[0]
|
||||
}
|
||||
wantGraph := append([]string{wantType}, want...)
|
||||
if len(want) > 0 {
|
||||
wantGraph = wantGraph[1:]
|
||||
}
|
||||
wantGraph = append(wantGraph, fmt.Sprintf("legacy_entity:%d", id))
|
||||
for _, term := range []string{"Trust", "Learning"} {
|
||||
graph, err := r.SearchNode(ctx, ns, []string{term})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var found *types.GraphNode
|
||||
for _, node := range graph.Node {
|
||||
if node.Name == "Trust" {
|
||||
found = node
|
||||
}
|
||||
}
|
||||
if found == nil || !slices.Equal(found.Attributes, wantGraph) {
|
||||
t.Fatalf("%s returned unsupported native attributes: got=%#v want=%v", term, found, wantGraph)
|
||||
}
|
||||
}
|
||||
for _, entityID := range []int64{0, id} {
|
||||
page, err := r.Browse(ctx, "kb", "", 0, entityID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, node := range page.Nodes {
|
||||
if node.ID != id {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
var payload struct {
|
||||
Attributes []string `json:"attributes"`
|
||||
Type string `json:"entity_type"`
|
||||
}
|
||||
if err := json.Unmarshal(node.Payload, &payload); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Equal(payload.Attributes, want) || payload.Type != wantType {
|
||||
t.Fatalf("browse returned unsupported native attributes: %#v; want %v / %s", payload, want, wantType)
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("browse lost the entity's surviving bare reference")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeReadersRevalidateSharedAttributeSources(t *testing.T) {
|
||||
for name, invalidate := range map[string]string{
|
||||
"changed source": `UPDATE chunks SET content=content || ' changed' WHERE id='chunk' AND tenant_id=7`,
|
||||
"disabled source": `UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`,
|
||||
"deleted source": `UPDATE chunks SET deleted_at=now() WHERE id='chunk' AND tenant_id=7`,
|
||||
"moved source document": `UPDATE chunks SET knowledge_id='doc2' WHERE id='chunk' AND tenant_id=7`,
|
||||
"wrong evidence hash": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{source_content_hash}',to_jsonb(repeat('a',64))) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong row hash": `UPDATE native_graph_records SET source_snapshot_sha256=repeat('a',64) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"inactive evidence": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{status}','"inactive"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong origin": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{origin}','"foreign"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong tenant": `UPDATE native_graph_records SET tenant_id=8 WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong knowledge base": `UPDATE native_graph_records SET knowledge_base_id='other-kb' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong document": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{native_knowledge_id}','"doc2"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong mapping document": `UPDATE native_graph_sources SET native_knowledge_id='doc2' WHERE source_chunk_id='native:chunk'`,
|
||||
"wrong mapping hash": `UPDATE native_graph_sources SET source_content_hash=repeat('a',64) WHERE source_chunk_id='native:chunk'`,
|
||||
"inactive mapping": `UPDATE native_graph_sources SET mapping_state='stale' WHERE source_chunk_id='native:chunk'`,
|
||||
"missing old snapshot": `UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"null snapshot": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}','null'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk' AND payload->>'subject_type'='entity'`,
|
||||
"null attribute": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}','[null]'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk' AND payload->>'subject_type'='entity'`,
|
||||
"conflicting snapshot": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}','["contradiction"]'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:sibling' AND payload->>'subject_type'='entity'`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
owner, bare := testGraph("chunk"), testGraph("sibling")
|
||||
owner.Node[0].Attributes = []string{"restriction", "Cannot be taken orally"}
|
||||
bare.Node[0].Attributes = nil
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{owner, bare}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertReaderTrustAttributes(t, r, ctx, owner.Node[0].Attributes)
|
||||
if err := tx.Exec(invalidate).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before := nativeGraphWriteState(t, tx)
|
||||
assertReaderTrustAttributes(t, r, ctx, nil)
|
||||
if after := nativeGraphWriteState(t, tx); after != before {
|
||||
t.Fatal("reader projection changed stored graph records or mappings")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeReadersKeepVerifiedAttributesBeyondChunkCap(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks SELECT id,7,'kb','doc','Another source about trust and learning.',NULL,true FROM unnest(ARRAY['aa','ab','ac','ad']) id`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
owner := testGraph("chunk")
|
||||
owner.Node[0].Attributes = []string{"restriction", "Cannot be taken orally"}
|
||||
graphs := []*types.GraphData{owner}
|
||||
for _, id := range []string{"aa", "ab", "ac", "ad"} {
|
||||
bare := testGraph(id)
|
||||
bare.Node[0].Attributes = nil
|
||||
graphs = append(graphs, bare)
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, graphs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// An old bare source cannot erase attributes proved by an explicit source.
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:aa'`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertReaderTrustAttributes(t, r, ctx, owner.Node[0].Attributes)
|
||||
graph, err := r.SearchNode(ctx, ns, []string{"Trust"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, node := range graph.Node {
|
||||
if len(node.Chunks) != 4 || (node.Name == "Trust" && !slices.Contains(node.Chunks, "chunk")) {
|
||||
t.Fatalf("attribute owner was omitted from the bounded references: %#v", node)
|
||||
}
|
||||
}
|
||||
if err := tx.Exec(`UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertReaderTrustAttributes(t, r, ctx, nil)
|
||||
}
|
||||
|
||||
func TestNativeReadersRespectCopiedSourceMappings(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
graph := testGraph("chunk")
|
||||
graph.Node[0].Attributes = []string{"restriction", "Cannot be taken orally"}
|
||||
if err := r.AddGraph(ctx, types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, sql := range []string{
|
||||
`INSERT INTO knowledges VALUES('copy-doc',7,'other-kb','copy.md',NULL)`,
|
||||
`INSERT INTO chunks VALUES('copy-chunk',7,'other-kb','copy-doc','Evidence about trust and learning.',NULL,true)`,
|
||||
`INSERT INTO native_graph_records SELECT tenant_id,'other-kb',source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256,imported_at FROM native_graph_records WHERE tenant_id=7 AND knowledge_base_id='kb'`,
|
||||
`INSERT INTO native_graph_sources SELECT tenant_id,'other-kb',source_chunk_id,source_file,source_content_hash,excerpt,'copy-doc','copy-chunk',native_content_sha256,mapping_state FROM native_graph_sources WHERE tenant_id=7 AND knowledge_base_id='kb'`,
|
||||
`UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`,
|
||||
} {
|
||||
if err := tx.Exec(sql).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "other-kb", Knowledge: "copy-doc"}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
found, err := r.SearchNode(ctx, ns, []string{"Trust"})
|
||||
if err != nil || len(found.Node) != 2 {
|
||||
t.Fatalf("copied native graph was not retrieved: %#v %v", found, err)
|
||||
}
|
||||
for _, node := range found.Node {
|
||||
if !slices.Equal(node.Chunks, []string{"copy-chunk"}) {
|
||||
t.Fatalf("copied graph retained an original chunk reference: %#v", node)
|
||||
}
|
||||
if node.Name == "Trust" && !slices.Equal(node.Attributes, []string{"restriction", "Cannot be taken orally", fmt.Sprintf("legacy_entity:%d", id)}) {
|
||||
t.Fatalf("copied source snapshot lost its attributes: %#v", node)
|
||||
}
|
||||
}
|
||||
wrong, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: "other-kb", Knowledge: "doc"}, []string{"Trust"})
|
||||
if err != nil || len(wrong.Node) != 0 {
|
||||
t.Fatalf("copied graph leaked through the original document scope: %#v %v", wrong, err)
|
||||
}
|
||||
page, err := r.Browse(ctx, "other-kb", "", 0, id)
|
||||
if err != nil || len(page.Nodes) != 1 {
|
||||
t.Fatalf("copied entity was not browsable: %#v %v", page, err)
|
||||
}
|
||||
var payload struct {
|
||||
Attributes []string `json:"attributes"`
|
||||
}
|
||||
if err := json.Unmarshal(page.Nodes[0].Payload, &payload); err != nil || !slices.Equal(payload.Attributes, graph.Node[0].Attributes) {
|
||||
t.Fatalf("copied browse attributes were not revalidated: %#v %v", payload, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Fixture mutations are confined to a rolled-back transaction.
|
||||
func TestAcceptancePreservedGraph(t *testing.T) {
|
||||
dsn := os.Getenv("GRAPH_TEST_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("GRAPH_TEST_DSN is required for acceptance integration")
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tx := db.Begin()
|
||||
defer tx.Rollback()
|
||||
if tx.Error != nil {
|
||||
t.Fatal(tx.Error)
|
||||
}
|
||||
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(10059))
|
||||
r := NewRepository(tx, nil)
|
||||
var fixture struct {
|
||||
KB string
|
||||
Entity int64
|
||||
Name string
|
||||
Chunk string
|
||||
Knowledge string
|
||||
}
|
||||
err = tx.Raw(`SELECT r.knowledge_base_id AS kb,r.legacy_id AS entity,r.payload->>'canonical_name' AS name,s.native_chunk_id AS chunk,s.native_knowledge_id AS knowledge FROM native_graph_records r JOIN native_graph_records e ON e.tenant_id=r.tenant_id AND e.knowledge_base_id=r.knowledge_base_id AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_type'='entity' AND e.payload->>'subject_id'=r.legacy_id::text JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id' WHERE r.tenant_id=10059 AND r.source_table='knowledgeentity' AND r.knowledge_base_id='16c3e59c-4bb6-41cf-ae0a-6d0941b3d032' AND s.native_chunk_id IS NOT NULL ORDER BY r.legacy_id LIMIT 1`).Scan(&fixture).Error
|
||||
if err != nil || fixture.Chunk == "" {
|
||||
t.Fatalf("missing mapped fixture: %v", err)
|
||||
}
|
||||
overview, err := r.Browse(ctx, fixture.KB, "", 0, fixture.Entity)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var preservedEntities, preservedMappings int64
|
||||
if err := tx.Raw(`SELECT count(*) FROM native_graph_records WHERE tenant_id=10059 AND knowledge_base_id=? AND source_table='knowledgeentity' AND payload->>'origin' IS DISTINCT FROM 'native_extraction_v1'`, fixture.KB).Scan(&preservedEntities).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Raw(`SELECT count(*) FROM native_graph_sources WHERE tenant_id=10059 AND knowledge_base_id=? AND mapping_state<>'native_extraction' AND native_chunk_id IS NOT NULL`, fixture.KB).Scan(&preservedMappings).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(overview.Nodes) != 1 || len(overview.Evidence) == 0 || preservedEntities != 2 || preservedMappings != 3 || overview.Counts["knowledgeentity"] < preservedEntities || overview.MappedSources < preservedMappings {
|
||||
t.Fatalf("preserved graph changed: entities=%d mappings=%d counts=%v", preservedEntities, preservedMappings, overview.Counts)
|
||||
}
|
||||
found := false
|
||||
for _, e := range overview.Evidence {
|
||||
if e.NativeChunkID == fixture.Chunk {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("mapped evidence was not returned")
|
||||
}
|
||||
graph, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: fixture.KB, Knowledge: fixture.Knowledge}, []string{fixture.Name})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(graph.Node) == 0 {
|
||||
t.Fatal("native graph retrieval produced no nodes")
|
||||
}
|
||||
for _, n := range graph.Node {
|
||||
for _, c := range n.Chunks {
|
||||
if strings.HasPrefix(c, "chunk-") {
|
||||
t.Fatal("legacy chunk ID leaked into native retrieval")
|
||||
}
|
||||
}
|
||||
}
|
||||
other := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(999999999))
|
||||
denied, err := r.Browse(other, fixture.KB, "", 0, fixture.Entity)
|
||||
if err != nil || len(denied.Nodes) != 0 || len(denied.Evidence) != 0 {
|
||||
t.Fatal("cross-tenant graph leak", err)
|
||||
}
|
||||
wrongFile, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: fixture.KB, Knowledge: "other-file"}, []string{fixture.Name})
|
||||
if err != nil || len(wrongFile.Node) != 0 {
|
||||
t.Fatal("file scope was ignored", err)
|
||||
}
|
||||
for _, change := range []struct {
|
||||
name, sql, id string
|
||||
}{
|
||||
{"edited chunk", "UPDATE chunks SET content=content || ' stale-check' WHERE id=? AND tenant_id=10059", fixture.Chunk},
|
||||
{"disabled chunk", "UPDATE chunks SET is_enabled=false WHERE id=? AND tenant_id=10059", fixture.Chunk},
|
||||
{"deleted chunk", "UPDATE chunks SET deleted_at=now() WHERE id=? AND tenant_id=10059", fixture.Chunk},
|
||||
{"deleted document", "UPDATE knowledges SET deleted_at=now() WHERE id=? AND tenant_id=10059", fixture.Knowledge},
|
||||
{"deleted knowledge base", "UPDATE knowledge_bases SET deleted_at=now() WHERE id=? AND tenant_id=10059", fixture.KB},
|
||||
} {
|
||||
t.Run(change.name, func(t *testing.T) {
|
||||
if err := tx.SavePoint("resource_check").Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer tx.RollbackTo("resource_check")
|
||||
if err := tx.Exec(change.sql, change.id).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stale, err := r.Browse(ctx, fixture.KB, "", 0, fixture.Entity)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, e := range stale.Evidence {
|
||||
if e.NativeChunkID == fixture.Chunk {
|
||||
t.Fatal("unavailable chunk remained linked")
|
||||
}
|
||||
}
|
||||
staleGraph, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: fixture.KB}, []string{fixture.Name})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, n := range staleGraph.Node {
|
||||
for _, c := range n.Chunks {
|
||||
if c == fixture.Chunk {
|
||||
t.Fatal("unavailable chunk leaked into retrieval")
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceLargestGraph(t *testing.T) {
|
||||
dsn := os.Getenv("GRAPH_TEST_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("GRAPH_TEST_DSN is required for acceptance integration")
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(10059))
|
||||
kb := "bb86b44a-576b-4bf8-a6bf-8f1baca3a845"
|
||||
r := NewRepository(db, nil)
|
||||
started := time.Now()
|
||||
overview, err := r.Browse(ctx, kb, "", 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var preserved struct{ Sources, Mapped int64 }
|
||||
if err := db.Raw(`SELECT count(*) AS sources,count(native_chunk_id) AS mapped FROM native_graph_sources WHERE tenant_id=10059 AND knowledge_base_id=? AND mapping_state<>'native_extraction'`, kb).Scan(&preserved).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(overview.Nodes) != 50 || preserved.Sources != 66204 || preserved.Mapped != 7000 || overview.Sources < preserved.Sources || overview.MappedSources < preserved.Mapped {
|
||||
t.Fatalf("unexpected largest graph overview: preserved=%+v sources=%d mapped=%d", preserved, overview.Sources, overview.MappedSources)
|
||||
}
|
||||
t.Logf("largest KB overview: %s", time.Since(started))
|
||||
started = time.Now()
|
||||
graph, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: kb}, []string{"销售", "客户", "信任", "心法"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(graph.Node) == 0 || len(graph.Node) > 32 || len(graph.Relation) > 32 {
|
||||
t.Fatalf("invalid bounded retrieval: %d nodes, %d relations", len(graph.Node), len(graph.Relation))
|
||||
}
|
||||
for _, node := range graph.Node {
|
||||
if len(node.Chunks) == 0 || len(node.Chunks) > 4 {
|
||||
t.Fatalf("invalid chunk cap: %#v", node)
|
||||
}
|
||||
}
|
||||
t.Logf("largest KB graph retrieval: %s; %d nodes, %d relations", time.Since(started), len(graph.Node), len(graph.Relation))
|
||||
|
||||
var relationFixture struct {
|
||||
KB, From, To, Relation string
|
||||
}
|
||||
err = db.Raw(`SELECT r.knowledge_base_id AS kb,a.payload->>'canonical_name' AS "from",b.payload->>'canonical_name' AS "to",r.payload->>'relation_type' AS relation FROM native_graph_records r JOIN native_graph_records a ON a.tenant_id=r.tenant_id AND a.knowledge_base_id=r.knowledge_base_id AND a.source_table='knowledgeentity' AND a.legacy_id=(r.payload->>'from_entity_id')::bigint JOIN native_graph_records b ON b.tenant_id=r.tenant_id AND b.knowledge_base_id=r.knowledge_base_id AND b.source_table='knowledgeentity' AND b.legacy_id=(r.payload->>'to_entity_id')::bigint WHERE r.tenant_id=10059 AND r.source_table='knowledgerelation' AND r.legacy_id=1`).Scan(&relationFixture).Error
|
||||
if err != nil || relationFixture.KB == "" {
|
||||
t.Fatalf("missing relation fixture: %v", err)
|
||||
}
|
||||
expanded, err := r.SearchNode(ctx, types.NameSpace{KnowledgeBase: relationFixture.KB}, []string{relationFixture.From})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, relation := range expanded.Relation {
|
||||
if relation.Node1 == relationFixture.From && relation.Node2 == relationFixture.To && relation.Type == relationFixture.Relation {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("verified relation neighbor was not expanded")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,522 @@
|
||||
package legacygraph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const nativeOrigin = "native_extraction_v1"
|
||||
|
||||
type sourceChunk struct {
|
||||
ID, KnowledgeID, KnowledgeBaseID, Content, FileName string
|
||||
}
|
||||
|
||||
func contentHash(content string) string {
|
||||
h := sha256.Sum256([]byte(content))
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
// IDs stay exactly representable in the browser, outside historical IDs. The
|
||||
// full identity is also checked on conflict, so a hash collision fails closed.
|
||||
func nativeID(parts ...string) (int64, string) {
|
||||
identity, _ := json.Marshal(parts)
|
||||
h := sha256.Sum256(identity)
|
||||
return int64(binary.BigEndian.Uint64(h[:8])&((1<<52)-1)) | 1<<52, hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
func normalizedName(name string) string { return strings.ToLower(strings.TrimSpace(name)) }
|
||||
|
||||
func sourceScopedEntity(name string) bool {
|
||||
name = strings.Map(func(r rune) rune {
|
||||
if unicode.IsSpace(r) {
|
||||
return -1
|
||||
}
|
||||
return r
|
||||
}, normalizedName(name))
|
||||
switch name {
|
||||
case "本品", "该产品", "本产品", "产品", "本商品", "此产品", "本公司", "该公司", "本文", "本方法", "该方法", "thisproduct", "theproduct":
|
||||
return true
|
||||
default:
|
||||
// The extractor uses these names to separate unresolved local products.
|
||||
return strings.HasSuffix(name, "未名产品") || strings.HasPrefix(name, "本品(")
|
||||
}
|
||||
}
|
||||
|
||||
func nativeEntityIdentity(document string, node *types.GraphNode) (int64, string, error) {
|
||||
key := normalizedName(node.Name)
|
||||
parts := []string{document, "entity", key}
|
||||
if sourceScopedEntity(key) {
|
||||
// An unnamed referent is not a document-wide entity. Multiple sources
|
||||
// require the extractor to supply an unambiguous, specifically named node.
|
||||
if len(node.Chunks) == 0 || node.Chunks[0] == "" {
|
||||
return 0, "", fmt.Errorf("generic graph entity requires one source chunk")
|
||||
}
|
||||
for _, chunk := range node.Chunks {
|
||||
if chunk != node.Chunks[0] {
|
||||
return 0, "", fmt.Errorf("generic graph entity cannot merge distinct source chunks")
|
||||
}
|
||||
}
|
||||
parts = append(parts, "source", node.Chunks[0])
|
||||
}
|
||||
id, identity := nativeID(parts...)
|
||||
return id, identity, nil
|
||||
}
|
||||
|
||||
// HasChunkGraph accepts only current checkpoints with extracted relations.
|
||||
// Legacy empty results and excerpt mappings require coverage review.
|
||||
func (r *Repository) HasChunkGraph(ctx context.Context, chunk *types.Chunk) (bool, error) {
|
||||
tenant, ok := types.TenantIDFromContext(ctx)
|
||||
if !ok || tenant == 0 || chunk == nil || chunk.TenantID != tenant || r.db == nil {
|
||||
return false, fmt.Errorf("graph checkpoint requires source tenant and chunk")
|
||||
}
|
||||
var count int64
|
||||
err := r.db.WithContext(ctx).Table("native_graph_records").Where("tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphchunkstate' AND payload->>'origin'=? AND payload->>'native_chunk_id'=? AND payload->>'native_knowledge_id'=? AND payload->>'source_content_hash'=? AND payload->>'extraction_status'='complete' AND CAST(COALESCE(payload->>'relation_count','0') AS INTEGER)>0", tenant, chunk.KnowledgeBaseID, nativeOrigin, chunk.ID, chunk.KnowledgeID, contentHash(chunk.Content)).Count(&count).Error
|
||||
return count > 0, err
|
||||
}
|
||||
|
||||
// AddChunkGraph rejects unreviewed empty coverage and results from changed sources.
|
||||
func (r *Repository) AddChunkGraph(ctx context.Context, ns types.NameSpace, chunk *types.Chunk, graph *types.GraphData) error {
|
||||
if chunk == nil || chunk.ID == "" || graph == nil {
|
||||
return fmt.Errorf("graph extraction requires its source chunk and model result")
|
||||
}
|
||||
if len(graph.Relation) == 0 {
|
||||
return fmt.Errorf("graph extraction without relations requires coverage review")
|
||||
}
|
||||
return r.addNativeGraphs(ctx, ns, []*types.GraphData{graph}, map[string]string{chunk.ID: contentHash(chunk.Content)})
|
||||
}
|
||||
|
||||
func (r *Repository) addNativeGraphs(ctx context.Context, ns types.NameSpace, graphs []*types.GraphData, expected map[string]string) error {
|
||||
tenant, ok := types.TenantIDFromContext(ctx)
|
||||
if !ok || tenant == 0 || ns.KnowledgeBase == "" || ns.Knowledge == "" || r.db == nil {
|
||||
return fmt.Errorf("graph writes require tenant, knowledge base, document and database")
|
||||
}
|
||||
ids := map[string]bool{}
|
||||
for id := range expected {
|
||||
ids[id] = true
|
||||
}
|
||||
for _, graph := range graphs {
|
||||
if graph == nil {
|
||||
return fmt.Errorf("nil graph extraction")
|
||||
}
|
||||
for _, node := range graph.Node {
|
||||
if node == nil || normalizedName(node.Name) == "" || len(node.Chunks) == 0 {
|
||||
return fmt.Errorf("extracted graph node requires name and source chunks")
|
||||
}
|
||||
for _, id := range node.Chunks {
|
||||
if id == "" || (expected != nil && expected[id] == "") {
|
||||
return fmt.Errorf("graph node refers to a chunk outside extraction scope")
|
||||
}
|
||||
ids[id] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
return fmt.Errorf("graph extraction has no source chunks")
|
||||
}
|
||||
chunkIDs := make([]string, 0, len(ids))
|
||||
for id := range ids {
|
||||
chunkIDs = append(chunkIDs, id)
|
||||
}
|
||||
sort.Strings(chunkIDs)
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := lockGraph(tx, tenant, ns.KnowledgeBase); err != nil {
|
||||
return err
|
||||
}
|
||||
var chunks []sourceChunk
|
||||
if err := tx.Raw(`SELECT c.id,c.knowledge_id,c.knowledge_base_id,c.content,k.file_name FROM chunks c
|
||||
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
|
||||
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
|
||||
WHERE c.tenant_id=? AND c.knowledge_base_id=? AND c.knowledge_id=? AND c.id IN ? AND c.deleted_at IS NULL AND c.is_enabled=true
|
||||
ORDER BY c.id FOR SHARE OF c,k,kb`, tenant, ns.KnowledgeBase, ns.Knowledge, chunkIDs).Scan(&chunks).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if len(chunks) != len(chunkIDs) {
|
||||
return fmt.Errorf("graph source chunks are unavailable or outside document/tenant scope")
|
||||
}
|
||||
byID := map[string]sourceChunk{}
|
||||
for _, chunk := range chunks {
|
||||
if expected != nil && expected[chunk.ID] != contentHash(chunk.Content) {
|
||||
return fmt.Errorf("graph source chunk %s changed during extraction", chunk.ID)
|
||||
}
|
||||
byID[chunk.ID] = chunk
|
||||
}
|
||||
// Replace only this extraction's evidence. Other documents, chunks and
|
||||
// archived migration records are untouched, including on retries.
|
||||
sources := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
source := "native:" + chunk.ID
|
||||
sources = append(sources, source)
|
||||
excerpt := []rune(chunk.Content)
|
||||
if len(excerpt) > 2000 {
|
||||
excerpt = excerpt[:2000]
|
||||
}
|
||||
if err := tx.Exec(`INSERT INTO native_graph_sources(tenant_id,knowledge_base_id,source_chunk_id,source_file,source_content_hash,excerpt,native_knowledge_id,native_chunk_id,native_content_sha256,mapping_state)
|
||||
VALUES(?,?,?,?,?,?,?,?,?,'native_extraction') ON CONFLICT(tenant_id,knowledge_base_id,source_chunk_id) DO UPDATE SET source_file=excluded.source_file,source_content_hash=excluded.source_content_hash,excerpt=excluded.excerpt,native_knowledge_id=excluded.native_knowledge_id,native_chunk_id=excluded.native_chunk_id,native_content_sha256=excluded.native_content_sha256,mapping_state=excluded.mapping_state`, tenant, ns.KnowledgeBase, source, chunk.FileName, contentHash(chunk.Content), string(excerpt), ns.Knowledge, chunk.ID, contentHash(chunk.Content)).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
type pendingEntity struct {
|
||||
identity string
|
||||
payload map[string]any
|
||||
}
|
||||
entities := map[int64]pendingEntity{}
|
||||
// Removed nodes also lose their source's attributes, even when a bare
|
||||
// reference from another chunk keeps the entity alive.
|
||||
var replacedEntities []Record
|
||||
if err := tx.Raw(`SELECT r.legacy_id,r.payload FROM native_graph_records r
|
||||
WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.source_table='knowledgeentity'
|
||||
AND r.payload->>'origin'=? AND r.payload->>'native_knowledge_id'=?
|
||||
AND EXISTS (SELECT 1 FROM native_graph_records e WHERE e.tenant_id=r.tenant_id AND e.knowledge_base_id=r.knowledge_base_id
|
||||
AND e.source_table='knowledgegraphevidence' AND e.payload->>'origin'=?
|
||||
AND e.payload->>'subject_type'='entity' AND e.payload->>'subject_id'=r.legacy_id::text AND e.payload->>'source_chunk_id' IN ?)`, tenant, ns.KnowledgeBase, nativeOrigin, ns.Knowledge, nativeOrigin, sources).Scan(&replacedEntities).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, record := range replacedEntities {
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(record.Payload, &payload); err != nil {
|
||||
return err
|
||||
}
|
||||
identity, _ := payload["native_identity"].(string)
|
||||
entities[record.ID] = pendingEntity{identity, payload}
|
||||
}
|
||||
if err := tx.Exec(`DELETE FROM native_graph_records WHERE tenant_id=? AND knowledge_base_id=? AND source_table='knowledgegraphevidence' AND payload->>'origin'=? AND payload->>'source_chunk_id' IN ?`, tenant, ns.KnowledgeBase, nativeOrigin, sources).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
counts := map[string]map[string]int{}
|
||||
entitySources := map[int64]map[string][]string{}
|
||||
for _, id := range chunkIDs {
|
||||
counts[id] = map[string]int{"entity": 0, "relation": 0}
|
||||
}
|
||||
for _, graph := range graphs {
|
||||
nodes := map[string]*types.GraphNode{}
|
||||
entityIDs := map[string]int64{}
|
||||
for _, node := range graph.Node {
|
||||
key := normalizedName(node.Name)
|
||||
id, identity, err := nativeEntityIdentity(ns.Knowledge, node)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if priorID, exists := entityIDs[key]; exists && priorID != id {
|
||||
return fmt.Errorf("generic graph entity is ambiguous within extraction")
|
||||
}
|
||||
name := strings.TrimSpace(node.Name)
|
||||
payload := map[string]any{"canonical_name": name, "entity_type": "concept", "attributes": node.Attributes}
|
||||
if len(node.Attributes) > 0 {
|
||||
payload["entity_type"] = node.Attributes[0]
|
||||
}
|
||||
if prior, exists := entities[id]; exists && prior.identity != identity {
|
||||
return fmt.Errorf("native graph identity collision; preserved graph was not modified")
|
||||
}
|
||||
if _, exists := entitySources[id]; !exists {
|
||||
entities[id] = pendingEntity{identity, payload}
|
||||
entitySources[id] = map[string][]string{}
|
||||
} else if name < entities[id].payload["canonical_name"].(string) {
|
||||
entities[id].payload["canonical_name"] = name
|
||||
}
|
||||
if prior := nodes[key]; prior != nil {
|
||||
prior.Chunks = append(prior.Chunks, node.Chunks...)
|
||||
} else {
|
||||
copyNode := *node
|
||||
copyNode.Chunks = append([]string(nil), node.Chunks...)
|
||||
nodes[key] = ©Node
|
||||
}
|
||||
entityIDs[key] = id
|
||||
seen := map[string]bool{}
|
||||
for _, chunkID := range node.Chunks {
|
||||
if seen[chunkID] {
|
||||
continue
|
||||
}
|
||||
seen[chunkID] = true
|
||||
if attributes, exists := entitySources[id][chunkID]; exists {
|
||||
if !slices.Equal(attributes, node.Attributes) {
|
||||
return fmt.Errorf("duplicate native graph entity source attributes conflict; extraction requires review")
|
||||
}
|
||||
continue
|
||||
}
|
||||
entitySources[id][chunkID] = node.Attributes
|
||||
if err := putEvidence(tx, tenant, ns, "entity", id, byID[chunkID], "", node.Attributes); err != nil {
|
||||
return err
|
||||
}
|
||||
counts[chunkID]["entity"]++
|
||||
}
|
||||
}
|
||||
for _, rel := range graph.Relation {
|
||||
if rel == nil || strings.TrimSpace(rel.Type) == "" {
|
||||
return fmt.Errorf("invalid extracted relation")
|
||||
}
|
||||
from, to := normalizedName(rel.Node1), normalizedName(rel.Node2)
|
||||
if nodes[from] == nil || nodes[to] == nil {
|
||||
return fmt.Errorf("relation references an entity outside extraction")
|
||||
}
|
||||
shared := []string{}
|
||||
fromChunks := map[string]bool{}
|
||||
for _, id := range nodes[from].Chunks {
|
||||
fromChunks[id] = true
|
||||
}
|
||||
for _, id := range nodes[to].Chunks {
|
||||
if fromChunks[id] {
|
||||
shared = append(shared, id)
|
||||
delete(fromChunks, id)
|
||||
}
|
||||
}
|
||||
if len(shared) == 0 {
|
||||
return fmt.Errorf("relation has no shared source evidence")
|
||||
}
|
||||
if rel.Evidence != "" {
|
||||
grounded := shared[:0]
|
||||
for _, chunkID := range shared {
|
||||
if strings.Contains(byID[chunkID].Content, rel.Evidence) {
|
||||
grounded = append(grounded, chunkID)
|
||||
}
|
||||
}
|
||||
shared = grounded
|
||||
if len(shared) == 0 {
|
||||
return fmt.Errorf("relation quotation is absent from current source evidence")
|
||||
}
|
||||
}
|
||||
identityParts := []string{ns.Knowledge, "relation", from, to, normalizedName(rel.Type)}
|
||||
if sourceScopedEntity(from) || sourceScopedEntity(to) {
|
||||
identityParts = append(identityParts, "scoped_endpoints", fmt.Sprint(entityIDs[from]), fmt.Sprint(entityIDs[to]))
|
||||
}
|
||||
id, identity := nativeID(identityParts...)
|
||||
if err := putNativeRecord(tx, tenant, ns, "knowledgerelation", id, identity, map[string]any{"from_entity_id": entityIDs[from], "to_entity_id": entityIDs[to], "relation_type": strings.TrimSpace(rel.Type), "evidence": rel.Evidence}, ""); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, chunkID := range shared {
|
||||
if err := putEvidence(tx, tenant, ns, "relation", id, byID[chunkID], rel.Evidence); err != nil {
|
||||
return err
|
||||
}
|
||||
counts[chunkID]["relation"]++
|
||||
}
|
||||
}
|
||||
}
|
||||
// Reconcile after every source snapshot is present, so batch ordering
|
||||
// cannot decide which source owns the entity's retained attributes.
|
||||
entityIDs := make([]int64, 0, len(entities))
|
||||
for id := range entities {
|
||||
entityIDs = append(entityIDs, id)
|
||||
}
|
||||
slices.Sort(entityIDs)
|
||||
for _, id := range entityIDs {
|
||||
entity := entities[id]
|
||||
if err := putNativeRecord(tx, tenant, ns, "knowledgeentity", id, entity.identity, entity.payload, ""); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, chunk := range chunks {
|
||||
id, identity := nativeID(ns.Knowledge, "chunkstate", chunk.ID)
|
||||
status := "complete"
|
||||
if counts[chunk.ID]["relation"] == 0 {
|
||||
status = "needs_review"
|
||||
}
|
||||
if err := putNativeRecord(tx, tenant, ns, "knowledgegraphchunkstate", id, identity, map[string]any{"source_chunk_id": "native:" + chunk.ID, "native_chunk_id": chunk.ID, "source_content_hash": contentHash(chunk.Content), "extraction_status": status, "entity_count": counts[chunk.ID]["entity"], "relation_count": counts[chunk.ID]["relation"]}, contentHash(chunk.Content)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return pruneNativeOrphans(tx, tenant, ns)
|
||||
})
|
||||
}
|
||||
|
||||
func lockGraph(tx *gorm.DB, tenant uint64, kb string) error {
|
||||
return tx.Exec("SELECT pg_advisory_xact_lock(hashtextextended(?,0))", fmt.Sprintf("native-graph:%d:%s", tenant, kb)).Error
|
||||
}
|
||||
|
||||
func putNativeRecord(tx *gorm.DB, tenant uint64, ns types.NameSpace, table string, id int64, identity string, payload map[string]any, hash string) error {
|
||||
if table == "knowledgeentity" {
|
||||
if err := reconcileNativeEntityAttributes(tx, tenant, ns, id, identity, payload); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
payload["origin"], payload["native_identity"], payload["native_knowledge_id"], payload["status"] = nativeOrigin, identity, ns.Knowledge, "active"
|
||||
data, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if hash == "" {
|
||||
hash = contentHash(string(data))
|
||||
}
|
||||
result := tx.Exec(`INSERT INTO native_graph_records(tenant_id,knowledge_base_id,source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256)
|
||||
VALUES(?,?,?,?,0,?::jsonb,?) ON CONFLICT(tenant_id,knowledge_base_id,source_table,legacy_id) DO UPDATE SET payload=excluded.payload,source_snapshot_sha256=excluded.source_snapshot_sha256
|
||||
WHERE native_graph_records.payload->>'native_identity'=excluded.payload->>'native_identity' AND native_graph_records.payload->>'origin'=?`, tenant, ns.KnowledgeBase, table, id, string(data), hash, nativeOrigin)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return fmt.Errorf("native graph identity collision; preserved graph was not modified")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func reconcileNativeEntityAttributes(tx *gorm.DB, tenant uint64, ns types.NameSpace, id int64, identity string, payload map[string]any) error {
|
||||
var evidence []struct {
|
||||
Payload json.RawMessage
|
||||
}
|
||||
// Evidence for the replaced sources has already been rewritten with raw
|
||||
// node attributes. Only still-current sources may preserve shared facts.
|
||||
if err := tx.Raw(`SELECT e.payload FROM native_graph_records e
|
||||
JOIN native_graph_sources s ON s.tenant_id=e.tenant_id AND s.knowledge_base_id=e.knowledge_base_id AND s.source_chunk_id=e.payload->>'source_chunk_id'
|
||||
JOIN chunks c ON c.id=s.native_chunk_id AND c.tenant_id=s.tenant_id AND c.knowledge_base_id=s.knowledge_base_id
|
||||
AND c.deleted_at IS NULL AND c.is_enabled=true AND encode(sha256(convert_to(c.content,'UTF8')),'hex')=s.native_content_sha256
|
||||
JOIN knowledges k ON k.id=c.knowledge_id AND k.tenant_id=c.tenant_id AND k.knowledge_base_id=c.knowledge_base_id AND k.deleted_at IS NULL
|
||||
JOIN knowledge_bases kb ON kb.id=c.knowledge_base_id AND kb.tenant_id=c.tenant_id AND kb.deleted_at IS NULL
|
||||
WHERE e.tenant_id=? AND e.knowledge_base_id=? AND e.source_table='knowledgegraphevidence'
|
||||
AND e.payload->>'subject_type'='entity' AND e.payload->>'subject_id'=?
|
||||
AND e.payload->>'origin'=? AND e.payload->>'status'='active' AND e.payload->>'native_knowledge_id'=?
|
||||
AND c.knowledge_id=? AND s.native_knowledge_id=c.knowledge_id AND s.mapping_state='native_extraction'
|
||||
AND e.payload->>'source_content_hash'=s.native_content_sha256 AND s.source_content_hash=s.native_content_sha256
|
||||
AND e.source_snapshot_sha256=s.native_content_sha256
|
||||
ORDER BY c.id,e.legacy_id FOR SHARE OF c,k,kb`, tenant, ns.KnowledgeBase, fmt.Sprint(id), nativeOrigin, ns.Knowledge, ns.Knowledge).Scan(&evidence).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
attributes := []string{}
|
||||
unknown := false
|
||||
for _, row := range evidence {
|
||||
var source struct {
|
||||
Attributes json.RawMessage `json:"entity_attributes"`
|
||||
}
|
||||
if err := json.Unmarshal(row.Payload, &source); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(source.Attributes) == 0 {
|
||||
unknown = true
|
||||
continue
|
||||
}
|
||||
var snapshot []*string
|
||||
if err := json.Unmarshal(source.Attributes, &snapshot); err != nil {
|
||||
return fmt.Errorf("invalid native graph entity attribute evidence: %w", err)
|
||||
}
|
||||
if snapshot == nil {
|
||||
unknown = true
|
||||
continue
|
||||
}
|
||||
values := make([]string, len(snapshot))
|
||||
for i, value := range snapshot {
|
||||
if value == nil {
|
||||
return fmt.Errorf("invalid native graph entity attribute evidence: null attribute")
|
||||
}
|
||||
values[i] = *value
|
||||
}
|
||||
if len(values) == 0 {
|
||||
continue
|
||||
}
|
||||
if len(attributes) > 0 && !slices.Equal(attributes, values) {
|
||||
return fmt.Errorf("shared native graph entity attributes conflict; extraction requires review")
|
||||
}
|
||||
attributes = values
|
||||
}
|
||||
entityType := "concept"
|
||||
if len(attributes) > 0 {
|
||||
entityType = attributes[0]
|
||||
}
|
||||
if unknown {
|
||||
var existing Record
|
||||
result := tx.Raw(`SELECT payload FROM native_graph_records WHERE tenant_id=? AND knowledge_base_id=?
|
||||
AND source_table='knowledgeentity' AND legacy_id=? AND payload->>'origin'=? AND payload->>'native_identity'=?`, tenant, ns.KnowledgeBase, id, nativeOrigin, identity).Scan(&existing)
|
||||
if result.Error != nil {
|
||||
return result.Error
|
||||
}
|
||||
var prior struct {
|
||||
Attributes []string `json:"attributes"`
|
||||
Type string `json:"entity_type"`
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return fmt.Errorf("shared native graph entity attribute provenance is unavailable; extraction requires review")
|
||||
}
|
||||
if err := json.Unmarshal(existing.Payload, &prior); err != nil {
|
||||
return err
|
||||
}
|
||||
// Legacy evidence cannot prove which surviving source owns a fact.
|
||||
// Refuse changing its aggregate until those sources are re-extracted.
|
||||
if !slices.Equal(prior.Attributes, attributes) || prior.Type != entityType {
|
||||
return fmt.Errorf("shared native graph entity attribute provenance is unavailable; extraction requires review")
|
||||
}
|
||||
attributes = prior.Attributes
|
||||
}
|
||||
payload["attributes"], payload["entity_type"] = attributes, entityType
|
||||
return nil
|
||||
}
|
||||
|
||||
func putEvidence(tx *gorm.DB, tenant uint64, ns types.NameSpace, subject string, subjectID int64, chunk sourceChunk, quotation string, attributes ...[]string) error {
|
||||
id, identity := nativeID(ns.Knowledge, "evidence", subject, fmt.Sprint(subjectID), chunk.ID)
|
||||
excerpt := []rune(chunk.Content)
|
||||
if len(excerpt) > 2000 {
|
||||
excerpt = excerpt[:2000]
|
||||
}
|
||||
if quotation != "" {
|
||||
excerpt = []rune(quotation)
|
||||
}
|
||||
payload := map[string]any{"subject_type": subject, "subject_id": subjectID, "source_chunk_id": "native:" + chunk.ID, "source_file": chunk.FileName, "source_content_hash": contentHash(chunk.Content), "excerpt": string(excerpt)}
|
||||
if subject == "entity" {
|
||||
if len(attributes) != 1 {
|
||||
return fmt.Errorf("native graph entity evidence requires source attributes")
|
||||
}
|
||||
payload["entity_attributes"] = append([]string{}, attributes[0]...)
|
||||
}
|
||||
return putNativeRecord(tx, tenant, ns, "knowledgegraphevidence", id, identity, payload, contentHash(chunk.Content))
|
||||
}
|
||||
|
||||
func pruneNativeOrphans(tx *gorm.DB, tenant uint64, ns types.NameSpace) error {
|
||||
query := `DELETE FROM native_graph_records r WHERE r.tenant_id=? AND r.knowledge_base_id=? AND r.payload->>'origin'=? AND r.source_table IN ('knowledgeentity','knowledgerelation')`
|
||||
args := []any{tenant, ns.KnowledgeBase, nativeOrigin}
|
||||
if ns.Knowledge != "" {
|
||||
query += " AND r.payload->>'native_knowledge_id'=?"
|
||||
args = append(args, ns.Knowledge)
|
||||
}
|
||||
query += ` AND NOT EXISTS (SELECT 1 FROM native_graph_records e WHERE e.tenant_id=r.tenant_id AND e.knowledge_base_id=r.knowledge_base_id AND e.source_table='knowledgegraphevidence' AND e.payload->>'subject_id'=r.legacy_id::text AND e.payload->>'subject_type'=CASE r.source_table WHEN 'knowledgeentity' THEN 'entity' ELSE 'relation' END)`
|
||||
return tx.Exec(query, args...).Error
|
||||
}
|
||||
|
||||
func (r *Repository) deleteNativeGraphs(ctx context.Context, namespaces []types.NameSpace) error {
|
||||
tenant, ok := types.TenantIDFromContext(ctx)
|
||||
if !ok || tenant == 0 || r.db == nil {
|
||||
return fmt.Errorf("graph deletion requires tenant and database")
|
||||
}
|
||||
for _, ns := range namespaces {
|
||||
if ns.KnowledgeBase == "" {
|
||||
return fmt.Errorf("graph deletion requires knowledge base scope")
|
||||
}
|
||||
}
|
||||
for _, ns := range namespaces {
|
||||
if err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := lockGraph(tx, tenant, ns.KnowledgeBase); err != nil {
|
||||
return err
|
||||
}
|
||||
query := tx.Table("native_graph_records").Where("tenant_id=? AND knowledge_base_id=? AND payload->>'origin'=?", tenant, ns.KnowledgeBase, nativeOrigin)
|
||||
if ns.Knowledge != "" {
|
||||
// Version copies preserve original payloads but remap source
|
||||
// pointers. Scope deletions through those current pointers too.
|
||||
query = query.Where("source_table IN ('knowledgegraphevidence','knowledgegraphchunkstate') AND (payload->>'native_knowledge_id'=? OR payload->>'source_chunk_id' IN (SELECT source_chunk_id FROM native_graph_sources WHERE tenant_id=? AND knowledge_base_id=? AND native_knowledge_id=?))", ns.Knowledge, tenant, ns.KnowledgeBase, ns.Knowledge)
|
||||
}
|
||||
if err := query.Delete(&Record{}).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if ns.Knowledge != "" {
|
||||
if err := pruneNativeOrphans(tx, tenant, types.NameSpace{KnowledgeBase: ns.KnowledgeBase}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
query = tx.Table("native_graph_sources").Where("tenant_id=? AND knowledge_base_id=? AND source_chunk_id LIKE 'native:%'", tenant, ns.KnowledgeBase)
|
||||
if ns.Knowledge != "" {
|
||||
query = query.Where("native_knowledge_id=?", ns.Knowledge)
|
||||
}
|
||||
return query.Delete(&Record{}).Error
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if r.fallback != nil {
|
||||
return r.fallback.DelGraph(ctx, namespaces)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,657 @@
|
||||
package legacygraph
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestNativeIdentity(t *testing.T) {
|
||||
a, identity := nativeID("doc", "entity", "concept")
|
||||
b, _ := nativeID("other-doc", "entity", "concept")
|
||||
if a < 1<<52 || a > (1<<53)-1 || a == b || len(identity) != 64 {
|
||||
t.Fatal("invalid stable document-scoped browser-safe ID")
|
||||
}
|
||||
c, again := nativeID("doc", "entity", "concept")
|
||||
if a != c || identity != again {
|
||||
t.Fatal("unstable identity")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeGenericEntityIdentityIsSourceScoped(t *testing.T) {
|
||||
for _, name := range []string{"本品", "本商品", "前段未名产品", "后段未名产品", "本品(前段)", "前段 未名产品", "本 品 (前段)", "This Product"} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
first, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{"first"}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{"second"}})
|
||||
if err != nil || first == second {
|
||||
t.Fatal("unnamed products from separate chunks were merged", err)
|
||||
}
|
||||
for _, chunks := range [][]string{nil, {""}, {"first", "second"}} {
|
||||
if _, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: chunks}); err == nil {
|
||||
t.Fatalf("missing or ambiguous source accepted: %v", chunks)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, name := range []string{"Named product", "修复凝胶", "本品专用护理套装", "未名产品研究中心"} {
|
||||
for _, chunk := range []string{"first", "second"} {
|
||||
id, _, err := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{chunk}})
|
||||
expected, _ := nativeID("doc", "entity", normalizedName(name))
|
||||
if err != nil || id != expected {
|
||||
t.Fatal("specific product lost its stable document identity", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeGenericProductsDoNotShareRelations(t *testing.T) {
|
||||
for _, name := range []string{"本品", "本商品", "前段未名产品", "后段未名产品", "本品(前段)"} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Evidence about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
for _, chunkID := range []string{"chunk", "sibling"} {
|
||||
graph := testGraph(chunkID)
|
||||
graph.Node[0].Name = name
|
||||
graph.Relation[0].Node1 = name
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
overview, err := r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil || overview.Counts["knowledgeentity"] != 3 || overview.Counts["knowledgerelation"] != 2 {
|
||||
t.Fatalf("separate unnamed products conflated: %#v %v", overview, err)
|
||||
}
|
||||
for _, chunkID := range []string{"chunk", "sibling"} {
|
||||
id, _, _ := nativeEntityIdentity("doc", &types.GraphNode{Name: name, Chunks: []string{chunkID}})
|
||||
detail, err := r.Browse(ctx, "kb", "", 0, id)
|
||||
if err != nil || len(detail.Relations) != 1 || len(detail.Evidence) != 1 || detail.Evidence[0].NativeChunkID != chunkID {
|
||||
t.Fatalf("unnamed product points to another source: %#v %v", detail, err)
|
||||
}
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{{Node: []*types.GraphNode{{Name: name, Chunks: []string{"chunk", "sibling"}}}}}); err == nil {
|
||||
t.Fatal("ambiguous multi-source write accepted")
|
||||
}
|
||||
after, err := r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil || after.Counts["knowledgerelation"] != 2 {
|
||||
t.Fatal("failed ambiguous write damaged committed graph", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func nativeWriteFixture(t *testing.T) (*gorm.DB, *Repository, context.Context) {
|
||||
t.Helper()
|
||||
dsn := os.Getenv("GRAPH_TEST_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("GRAPH_TEST_DSN required for temporary-table integration")
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { sqlDB.Close() })
|
||||
tx := db.Begin()
|
||||
if tx.Error != nil {
|
||||
t.Fatal(tx.Error)
|
||||
}
|
||||
t.Cleanup(func() { tx.Rollback() })
|
||||
// Temporary tables shadow live names on this transaction's connection;
|
||||
// tests never mutate application records or permanent schemas.
|
||||
for _, sql := range []string{
|
||||
`CREATE TEMP TABLE knowledge_bases(id text,tenant_id bigint,deleted_at timestamptz) ON COMMIT DROP`,
|
||||
`CREATE TEMP TABLE knowledges(id text,tenant_id bigint,knowledge_base_id text,file_name text,deleted_at timestamptz) ON COMMIT DROP`,
|
||||
`CREATE TEMP TABLE chunks(id text,tenant_id bigint,knowledge_base_id text,knowledge_id text,content text,deleted_at timestamptz,is_enabled boolean) ON COMMIT DROP`,
|
||||
`CREATE TEMP TABLE native_graph_records(tenant_id bigint,knowledge_base_id text,source_table text,legacy_id bigint,graph_version_id bigint,payload jsonb,source_snapshot_sha256 text,imported_at timestamptz DEFAULT now(),PRIMARY KEY(tenant_id,knowledge_base_id,source_table,legacy_id)) ON COMMIT DROP`,
|
||||
`CREATE TEMP TABLE native_graph_sources(tenant_id bigint,knowledge_base_id text,source_chunk_id text,source_file text,source_content_hash text,excerpt text,native_knowledge_id text,native_chunk_id text,native_content_sha256 text,mapping_state text,PRIMARY KEY(tenant_id,knowledge_base_id,source_chunk_id)) ON COMMIT DROP`,
|
||||
`INSERT INTO knowledge_bases VALUES('kb',7,NULL),('kb',8,NULL),('other-kb',7,NULL)`,
|
||||
`INSERT INTO knowledges VALUES('doc',7,'kb','source.md',NULL),('doc2',7,'kb','other.md',NULL),('doc',8,'kb','private.md',NULL)`,
|
||||
`INSERT INTO chunks VALUES('chunk',7,'kb','doc','Evidence about trust and learning.',NULL,true),('chunk2',7,'kb','doc2','Evidence about trust and learning.',NULL,true),('foreign',8,'kb','doc','Private source.',NULL,true)`,
|
||||
} {
|
||||
if err := tx.Exec(sql).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
|
||||
return tx, NewRepository(tx, nil), ctx
|
||||
}
|
||||
|
||||
func testGraph(chunk string) *types.GraphData {
|
||||
return &types.GraphData{Node: []*types.GraphNode{{Name: "Trust", Chunks: []string{chunk}, Attributes: []string{"concept", "Earned through reliable evidence"}}, {Name: "Learning", Chunks: []string{chunk}, Attributes: []string{"process"}}}, Relation: []*types.GraphRelation{{Node1: "Trust", Node2: "Learning", Type: "supports"}}}
|
||||
}
|
||||
|
||||
func TestNativeGraphWriteReplayBrowseAndIsolation(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Evidence about trust and learning."}
|
||||
for i := 0; i < 2; i++ {
|
||||
if err := r.AddChunkGraph(ctx, ns, chunk, testGraph("chunk")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
overview, err := r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if overview.Counts["knowledgeentity"] != 2 || overview.Counts["knowledgerelation"] != 1 || overview.Counts["knowledgegraphevidence"] != 3 || overview.Counts["knowledgegraphchunkstate"] != 1 || overview.MappedSources != 1 {
|
||||
t.Fatalf("not durable/idempotent: %#v", overview)
|
||||
}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
detail, err := r.Browse(ctx, "kb", "", 0, id)
|
||||
if err != nil || len(detail.Relations) != 1 || len(detail.Evidence) != 1 || detail.Evidence[0].NativeChunkID != "chunk" {
|
||||
t.Fatalf("detail missing actual evidence: %#v %v", detail, err)
|
||||
}
|
||||
found, err := r.SearchNode(ctx, ns, []string{"Trust"})
|
||||
if err != nil || len(found.Node) != 2 || len(found.Relation) != 1 {
|
||||
t.Fatalf("native RAG cannot find persisted graph: %#v %v", found, err)
|
||||
}
|
||||
complete, err := r.HasChunkGraph(ctx, chunk)
|
||||
if err != nil || !complete {
|
||||
t.Fatalf("missing checkpoint: %v", err)
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("foreign")}); err == nil {
|
||||
t.Fatal("cross-tenant source accepted")
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk2")}); err == nil {
|
||||
t.Fatal("cross-document source accepted")
|
||||
}
|
||||
changed := *chunk
|
||||
changed.Content = "old model input"
|
||||
if err := r.AddChunkGraph(ctx, ns, &changed, testGraph("chunk")); err == nil {
|
||||
t.Fatal("stale model output accepted")
|
||||
}
|
||||
bad := testGraph("chunk")
|
||||
bad.Relation[0].Node2 = "invented"
|
||||
if err := r.AddChunkGraph(ctx, ns, chunk, bad); err == nil {
|
||||
t.Fatal("ungrounded relation accepted")
|
||||
}
|
||||
after, err := r.Browse(ctx, "kb", "", 0, id)
|
||||
if err != nil || len(after.Evidence) != 1 {
|
||||
t.Fatal("failed extraction destroyed committed evidence", err)
|
||||
}
|
||||
if err := r.AddGraph(ctx, types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc2"}, []*types.GraphData{testGraph("chunk2")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Exec(`INSERT INTO native_graph_records VALUES(7,'kb','knowledgeentity',1,1,'{"canonical_name":"Archived","status":"active"}',repeat('a',64),now())`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.DelGraph(ctx, []types.NameSpace{ns}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
left, err := r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if left.Counts["knowledgeentity"] != 3 || left.Sources != 1 {
|
||||
t.Fatalf("deletion damaged other document/archive: %#v", left)
|
||||
}
|
||||
wrong := context.WithValue(ctx, types.TenantIDContextKey, uint64(8))
|
||||
if err := r.DelGraph(wrong, []types.NameSpace{{KnowledgeBase: "kb"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
left, err = r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil || left.Counts["knowledgeentity"] != 3 {
|
||||
t.Fatal("cross-tenant deletion", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeEmptyResultAndReextraction(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Evidence about trust and learning."}
|
||||
if err := r.AddChunkGraph(ctx, ns, chunk, testGraph("chunk")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.AddChunkGraph(ctx, ns, chunk, &types.GraphData{}); err == nil {
|
||||
t.Fatal("empty extraction accepted without review")
|
||||
}
|
||||
overview, err := r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil || overview.Counts["knowledgerelation"] == 0 || overview.Counts["knowledgegraphchunkstate"] != 1 {
|
||||
t.Fatalf("rejected result replaced prior graph: %#v %v", overview, err)
|
||||
}
|
||||
complete, err := r.HasChunkGraph(ctx, chunk)
|
||||
if err != nil || !complete {
|
||||
t.Fatal("prior verified graph was lost", err)
|
||||
}
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=jsonb_set(payload,'{relation_count}','0'::jsonb) WHERE source_table='knowledgegraphchunkstate'`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
complete, err = r.HasChunkGraph(ctx, chunk)
|
||||
if err != nil || complete {
|
||||
t.Fatal("legacy empty checkpoint was accepted", err)
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{{}}); err == nil {
|
||||
t.Fatal("empty extraction without source accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeRelationQuotationIsPersistedAndChecked(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
graph := testGraph("chunk")
|
||||
graph.Relation[0].Evidence = "trust and learning"
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var quote, excerpt string
|
||||
if err := tx.Raw(`SELECT payload->>'evidence' FROM native_graph_records WHERE source_table='knowledgerelation'`).Scan("e).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Raw(`SELECT payload->>'excerpt' FROM native_graph_records WHERE source_table='knowledgegraphevidence' AND payload->>'subject_type'='relation'`).Scan(&excerpt).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if quote != graph.Relation[0].Evidence || excerpt != quote {
|
||||
t.Fatalf("relation quotation not retained: %q %q", quote, excerpt)
|
||||
}
|
||||
graph.Relation[0].Evidence = "invented quotation"
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err == nil {
|
||||
t.Fatal("relation quotation absent from source accepted")
|
||||
}
|
||||
if err := tx.Raw(`SELECT payload->>'evidence' FROM native_graph_records WHERE source_table='knowledgerelation'`).Scan("e).Error; err != nil || quote != "trust and learning" {
|
||||
t.Fatal("invalid quotation damaged prior graph", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeRelationRequiresCurrentSourceEvidence(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source discussing trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
graph := testGraph("sibling")
|
||||
graph.Relation = nil
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{graph}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Exec(`UPDATE chunks SET content='The original relation is no longer supported.' WHERE id='chunk' AND tenant_id=7`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found, err := r.SearchNode(ctx, ns, []string{"Trust"})
|
||||
if err != nil || len(found.Node) != 2 || len(found.Relation) != 0 {
|
||||
t.Fatalf("stale relation remained available through separately supported entities: %#v %v", found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeCopiedVersionDeletionUsesRemappedSources(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, sql := range []string{
|
||||
`INSERT INTO knowledges VALUES('copy-doc',7,'other-kb','copy.md',NULL)`,
|
||||
`INSERT INTO chunks VALUES('copy-chunk',7,'other-kb','copy-doc','Evidence about trust and learning.',NULL,true)`,
|
||||
`INSERT INTO native_graph_records SELECT tenant_id,'other-kb',source_table,legacy_id,graph_version_id,payload,source_snapshot_sha256,imported_at FROM native_graph_records WHERE tenant_id=7 AND knowledge_base_id='kb'`,
|
||||
`INSERT INTO native_graph_sources SELECT tenant_id,'other-kb',source_chunk_id,source_file,source_content_hash,excerpt,'copy-doc','copy-chunk',native_content_sha256,mapping_state FROM native_graph_sources WHERE tenant_id=7 AND knowledge_base_id='kb'`,
|
||||
} {
|
||||
if err := tx.Exec(sql).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := r.DelGraph(ctx, []types.NameSpace{{KnowledgeBase: "other-kb", Knowledge: "copy-doc"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
copy, err := r.Browse(ctx, "other-kb", "", 0, 0)
|
||||
if err != nil || len(copy.Counts) != 0 || copy.Sources != 0 {
|
||||
t.Fatalf("copied graph was not deleted via its current document: %#v %v", copy, err)
|
||||
}
|
||||
original, err := r.Browse(ctx, "kb", "", 0, 0)
|
||||
if err != nil || original.Counts["knowledgeentity"] != 2 || original.Sources != 1 {
|
||||
t.Fatal("copy deletion changed original graph", err)
|
||||
}
|
||||
}
|
||||
|
||||
func nativeGraphWriteState(t *testing.T, tx *gorm.DB) string {
|
||||
t.Helper()
|
||||
var state string
|
||||
if err := tx.Raw(`SELECT jsonb_build_object(
|
||||
'records',COALESCE((SELECT jsonb_agg(to_jsonb(r) ORDER BY tenant_id,knowledge_base_id,source_table,legacy_id) FROM native_graph_records r),'[]'::jsonb),
|
||||
'sources',COALESCE((SELECT jsonb_agg(to_jsonb(s) ORDER BY tenant_id,knowledge_base_id,source_chunk_id) FROM native_graph_sources s),'[]'::jsonb))::text`).Scan(&state).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return state
|
||||
}
|
||||
|
||||
func assertNativeTrustAttributes(t *testing.T, tx *gorm.DB, want []string) string {
|
||||
t.Helper()
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
var record struct {
|
||||
Payload json.RawMessage
|
||||
Hash string
|
||||
State string
|
||||
}
|
||||
if err := tx.Raw(`SELECT payload,source_snapshot_sha256 AS hash,to_jsonb(r)::text AS state FROM native_graph_records r
|
||||
WHERE tenant_id=7 AND knowledge_base_id='kb' AND source_table='knowledgeentity' AND legacy_id=?`, id).Scan(&record).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var entity struct {
|
||||
Attributes []string `json:"attributes"`
|
||||
Type string `json:"entity_type"`
|
||||
}
|
||||
if err := json.Unmarshal(record.Payload, &entity); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantType := "concept"
|
||||
if len(want) > 0 {
|
||||
wantType = want[0]
|
||||
}
|
||||
if !slices.Equal(entity.Attributes, want) || entity.Type != wantType {
|
||||
t.Fatalf("unexpected aggregate attributes/type: %#v; want %v / %s", entity, want, wantType)
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(record.Payload, &payload); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil || record.Hash != contentHash(string(encoded)) {
|
||||
t.Fatalf("entity payload hash is stale: %s %v", record.Hash, err)
|
||||
}
|
||||
return record.State
|
||||
}
|
||||
|
||||
func TestNativeSharedEntityAttributesFollowTheirSources(t *testing.T) {
|
||||
for _, siblingHasAttributes := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("sibling_attributes_%t", siblingHasAttributes), func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
original := testGraph("chunk")
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before := assertNativeTrustAttributes(t, tx, original.Node[0].Attributes)
|
||||
sibling := testGraph("sibling")
|
||||
if !siblingHasAttributes {
|
||||
sibling.Node[0].Attributes = nil
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{sibling}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if after := assertNativeTrustAttributes(t, tx, original.Node[0].Attributes); after != before {
|
||||
t.Fatal("a second source changed the existing entity row")
|
||||
}
|
||||
var sourceAttributes string
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
if err := tx.Raw(`SELECT payload->'entity_attributes' FROM native_graph_records WHERE source_table='knowledgegraphevidence'
|
||||
AND payload->>'subject_type'='entity' AND payload->>'subject_id'=? AND payload->>'source_chunk_id'='native:sibling'`, fmt.Sprint(id)).Scan(&sourceAttributes).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !siblingHasAttributes && sourceAttributes != "[]" {
|
||||
t.Fatalf("bare source inherited aggregate attributes: %s", sourceAttributes)
|
||||
}
|
||||
original.Node[0].Attributes = nil
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, sibling.Node[0].Attributes)
|
||||
sibling.Node[0].Attributes = []string{}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{sibling}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, nil)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeSharedEntityAttributeConflictRollsBack(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before := nativeGraphWriteState(t, tx)
|
||||
conflicting := testGraph("sibling")
|
||||
conflicting.Node[0].Attributes = []string{"concept", "Never earned through reliable evidence"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{conflicting}); err == nil || !strings.Contains(err.Error(), "attributes conflict") {
|
||||
t.Fatalf("conflicting source did not require review: %v", err)
|
||||
}
|
||||
if after := nativeGraphWriteState(t, tx); after != before {
|
||||
t.Fatal("failed conflict changed records, evidence, mapping or checkpoint")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeSharedEntityAttributesAreRemovedWithTheirNode(t *testing.T) {
|
||||
for _, siblingHasAttributes := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("sibling_attributes_%t", siblingHasAttributes), func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
original, sibling := testGraph("chunk"), testGraph("sibling")
|
||||
if !siblingHasAttributes {
|
||||
sibling.Node[0].Attributes = nil
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{original, sibling}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
original.Node[0].Name = "Practice"
|
||||
original.Relation[0].Node1 = "Practice"
|
||||
chunk := &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Evidence about trust and learning."}
|
||||
if err := r.AddChunkGraph(ctx, ns, chunk, original); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, sibling.Node[0].Attributes)
|
||||
sibling.Node[0].Name = "Practice"
|
||||
sibling.Node[0].Attributes = original.Node[0].Attributes
|
||||
sibling.Relation[0].Node1 = "Practice"
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{sibling}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
var count int64
|
||||
if err := tx.Table("native_graph_records").Where("source_table='knowledgeentity' AND legacy_id=?", id).Count(&count).Error; err != nil || count != 0 {
|
||||
t.Fatalf("entity without surviving sources was not pruned: %d %v", count, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeLegacyEntityAttributeProvenanceRequiresReview(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence'
|
||||
AND payload->>'subject_type'='entity' AND payload->>'subject_id'=?`, fmt.Sprint(id)).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before := nativeGraphWriteState(t, tx)
|
||||
bare := testGraph("sibling")
|
||||
bare.Node[0].Attributes = nil
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err == nil || !strings.Contains(err.Error(), "provenance is unavailable") {
|
||||
t.Fatalf("legacy unknown provenance treated as known empty: %v", err)
|
||||
}
|
||||
if after := nativeGraphWriteState(t, tx); after != before {
|
||||
t.Fatal("unknown provenance failure changed prior graph")
|
||||
}
|
||||
// A matching explicit source can coexist with unknown older evidence.
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("sibling")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
original := testGraph("chunk")
|
||||
original.Node[0].Attributes = nil
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, testGraph("sibling").Node[0].Attributes)
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, nil)
|
||||
}
|
||||
|
||||
func TestNativeNullEntityAttributeEvidenceRequiresReview(t *testing.T) {
|
||||
for _, snapshot := range []string{"null", "[null]"} {
|
||||
t.Run(snapshot, func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
id, _ := nativeID("doc", "entity", "trust")
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=jsonb_set(payload,'{entity_attributes}',?::jsonb)
|
||||
WHERE source_table='knowledgegraphevidence' AND payload->>'subject_type'='entity' AND payload->>'subject_id'=?`, snapshot, fmt.Sprint(id)).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before := nativeGraphWriteState(t, tx)
|
||||
bare := testGraph("sibling")
|
||||
bare.Node[0].Attributes = nil
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err == nil {
|
||||
t.Fatal("null attribute evidence was accepted as known empty")
|
||||
}
|
||||
if after := nativeGraphWriteState(t, tx); after != before {
|
||||
t.Fatal("null attribute evidence failure changed prior graph")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeSharedEntityIgnoresInvalidAttributeEvidence(t *testing.T) {
|
||||
for name, invalidate := range map[string]string{
|
||||
"changed source": `UPDATE chunks SET content='Changed source.' WHERE id='chunk' AND tenant_id=7`,
|
||||
"disabled source": `UPDATE chunks SET is_enabled=false WHERE id='chunk' AND tenant_id=7`,
|
||||
"deleted source": `UPDATE chunks SET deleted_at=now() WHERE id='chunk' AND tenant_id=7`,
|
||||
"wrong evidence hash": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{source_content_hash}',to_jsonb(repeat('a',64))) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong row hash": `UPDATE native_graph_records SET source_snapshot_sha256=repeat('a',64) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"inactive evidence": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{status}','"inactive"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong origin": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{origin}','"foreign"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong tenant": `UPDATE native_graph_records SET tenant_id=8 WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong knowledge base": `UPDATE native_graph_records SET knowledge_base_id='other-kb' WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong document": `UPDATE native_graph_records SET payload=jsonb_set(payload,'{native_knowledge_id}','"doc2"'::jsonb) WHERE source_table='knowledgegraphevidence' AND payload->>'source_chunk_id'='native:chunk'`,
|
||||
"wrong mapping document": `UPDATE native_graph_sources SET native_knowledge_id='doc2' WHERE source_chunk_id='native:chunk'`,
|
||||
"inactive mapping": `UPDATE native_graph_sources SET mapping_state='stale' WHERE source_chunk_id='native:chunk'`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{testGraph("chunk")}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Exec(invalidate).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
bare := testGraph("sibling")
|
||||
bare.Node[0].Attributes = nil
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{bare}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, nil)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeEntityAttributeBatchOrderAndDuplicates(t *testing.T) {
|
||||
for _, reverse := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("reverse_%t", reverse), func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('sibling',7,'kb','doc','Another source about trust and learning.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
first, second := testGraph("chunk"), testGraph("sibling")
|
||||
first.Node = append(first.Node, &types.GraphNode{Name: " trust ", Chunks: []string{"chunk", "chunk"}, Attributes: first.Node[0].Attributes})
|
||||
second.Node[0].Name = "TRUST"
|
||||
second.Node[0].Attributes = nil
|
||||
graphs := []*types.GraphData{first, second}
|
||||
if reverse {
|
||||
slices.Reverse(graphs)
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, graphs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
entityBefore := assertNativeTrustAttributes(t, tx, first.Node[0].Attributes)
|
||||
slices.Reverse(graphs)
|
||||
if err := r.AddGraph(ctx, ns, graphs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if entityAfter := assertNativeTrustAttributes(t, tx, first.Node[0].Attributes); entityAfter != entityBefore {
|
||||
t.Fatal("batch order changed the entity payload or hash")
|
||||
}
|
||||
var count int
|
||||
if err := tx.Raw(`SELECT (payload->>'entity_count')::int FROM native_graph_records WHERE source_table='knowledgegraphchunkstate' AND payload->>'native_chunk_id'='chunk'`).Scan(&count).Error; err != nil || count != 2 {
|
||||
t.Fatalf("duplicate entity evidence inflated checkpoint: %d %v", count, err)
|
||||
}
|
||||
before := nativeGraphWriteState(t, tx)
|
||||
first.Node[2].Attributes = []string{"conflicting duplicate"}
|
||||
if err := r.AddGraph(ctx, ns, graphs); err == nil || !strings.Contains(err.Error(), "attributes conflict") {
|
||||
t.Fatalf("duplicate source attributes silently overwritten: %v", err)
|
||||
}
|
||||
if after := nativeGraphWriteState(t, tx); after != before {
|
||||
t.Fatal("failed duplicate source write changed the committed graph")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeLegacyEntityAttributeBatchOrder(t *testing.T) {
|
||||
for _, reverse := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("reverse_%t", reverse), func(t *testing.T) {
|
||||
tx, r, ctx := nativeWriteFixture(t)
|
||||
if err := tx.Exec(`INSERT INTO chunks VALUES('bare',7,'kb','doc','A bare reference to trust.',NULL,true),('explicit',7,'kb','doc','An explicit source about trust.',NULL,true)`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ns := types.NameSpace{KnowledgeBase: "kb", Knowledge: "doc"}
|
||||
original := testGraph("chunk")
|
||||
if err := r.AddGraph(ctx, ns, []*types.GraphData{original}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := tx.Exec(`UPDATE native_graph_records SET payload=payload-'entity_attributes' WHERE source_table='knowledgegraphevidence' AND payload->>'subject_type'='entity'`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
bare, explicit := testGraph("bare"), testGraph("explicit")
|
||||
bare.Node[0].Attributes = nil
|
||||
graphs := []*types.GraphData{bare, explicit}
|
||||
if reverse {
|
||||
slices.Reverse(graphs)
|
||||
}
|
||||
if err := r.AddGraph(ctx, ns, graphs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, original.Node[0].Attributes)
|
||||
original.Node[0].Attributes = nil
|
||||
explicit.Node[0].Attributes = nil
|
||||
graphs = append(graphs, original)
|
||||
if err := r.AddGraph(ctx, ns, graphs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertNativeTrustAttributes(t, tx, nil)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,304 @@
|
||||
package milvus
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"maps"
|
||||
"reflect"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
// operatorAnd is the "and" operator.
|
||||
operatorAnd = "and"
|
||||
|
||||
// operatorOr is the "or" operator.
|
||||
operatorOr = "or"
|
||||
|
||||
// operatorEqual is the "equal" operator.
|
||||
operatorEqual = "eq"
|
||||
|
||||
// operatorNotEqual is the "not equal" operator.
|
||||
operatorNotEqual = "ne"
|
||||
|
||||
// operatorGreaterThan is the "greater than" operator.
|
||||
operatorGreaterThan = "gt"
|
||||
|
||||
// operatorGreaterThanOrEqual is the "greater than or equal" operator.
|
||||
operatorGreaterThanOrEqual = "gte"
|
||||
|
||||
// operatorLessThan is the "less than" operator.
|
||||
operatorLessThan = "lt"
|
||||
|
||||
// operatorLessThanOrEqual is the "less than or equal" operator.
|
||||
operatorLessThanOrEqual = "lte"
|
||||
|
||||
// operatorIn is the "in" operator.
|
||||
operatorIn = "in"
|
||||
|
||||
// operatorNotIn is the "not in" operator.
|
||||
operatorNotIn = "not in"
|
||||
|
||||
// operatorLike is the "contains" operator.
|
||||
operatorLike = "like"
|
||||
|
||||
// operatorNotLike is the "not contains" operator.
|
||||
operatorNotLike = "not like"
|
||||
|
||||
// operatorBetween is the "between" operator.
|
||||
operatorBetween = "between"
|
||||
)
|
||||
|
||||
var comparisonOperators = map[string]string{
|
||||
operatorEqual: "==",
|
||||
operatorNotEqual: "!=",
|
||||
operatorGreaterThan: ">",
|
||||
operatorGreaterThanOrEqual: ">=",
|
||||
operatorLessThan: "<",
|
||||
operatorLessThanOrEqual: "<=",
|
||||
operatorLike: "like",
|
||||
operatorNotLike: "not like",
|
||||
}
|
||||
|
||||
type convertResult struct {
|
||||
exprStr string
|
||||
params map[string]any
|
||||
}
|
||||
|
||||
type filter struct{}
|
||||
|
||||
func (c *filter) Convert(cond *universalFilterCondition) (*convertResult, error) {
|
||||
var counter int
|
||||
return c.convertCondition(cond, &counter)
|
||||
}
|
||||
|
||||
func (c *filter) convertComparisonCondition(
|
||||
cond *universalFilterCondition,
|
||||
counter *int,
|
||||
) (*convertResult, error) {
|
||||
condField := cond.Field
|
||||
if condField == "" || cond.Value == nil {
|
||||
return nil, fmt.Errorf("milvus filter condition is nil")
|
||||
}
|
||||
operator, ok := comparisonOperators[cond.Operator]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("unsupported comparison operator: %s", cond.Operator)
|
||||
}
|
||||
|
||||
paramName := c.convertParamName(cond.Field, counter)
|
||||
return &convertResult{
|
||||
exprStr: fmt.Sprintf("%s %s {%s}", condField, operator, paramName),
|
||||
params: map[string]any{paramName: cond.Value},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *filter) convertLogicalCondition(
|
||||
cond *universalFilterCondition,
|
||||
counter *int,
|
||||
) (*convertResult, error) {
|
||||
if cond.Value == nil {
|
||||
return nil, fmt.Errorf("milvus filter condition is nil")
|
||||
}
|
||||
conds, ok := cond.Value.([]*universalFilterCondition)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("invalid logical condition value type")
|
||||
}
|
||||
|
||||
var condResult *convertResult
|
||||
for _, childCond := range conds {
|
||||
childRes, err := c.convertCondition(childCond, counter)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if childRes == nil || childRes.exprStr == "" {
|
||||
continue
|
||||
}
|
||||
if condResult == nil {
|
||||
condResult = childRes
|
||||
continue
|
||||
}
|
||||
|
||||
condResult.exprStr = fmt.Sprintf(
|
||||
"(%s) %s (%s)",
|
||||
condResult.exprStr,
|
||||
strings.ToLower(cond.Operator),
|
||||
childRes.exprStr,
|
||||
)
|
||||
maps.Copy(condResult.params, childRes.params)
|
||||
}
|
||||
|
||||
if condResult == nil {
|
||||
return nil, fmt.Errorf("empty logical condition")
|
||||
}
|
||||
return condResult, nil
|
||||
}
|
||||
|
||||
func (c *filter) convertCondition(
|
||||
cond *universalFilterCondition,
|
||||
counter *int,
|
||||
) (*convertResult, error) {
|
||||
if cond == nil {
|
||||
return nil, fmt.Errorf("milvus filter condition is nil")
|
||||
}
|
||||
switch cond.Operator {
|
||||
case operatorEqual, operatorNotEqual, operatorGreaterThan,
|
||||
operatorGreaterThanOrEqual, operatorLessThan,
|
||||
operatorLessThanOrEqual, operatorLike, operatorNotLike:
|
||||
return c.convertComparisonCondition(cond, counter)
|
||||
case operatorAnd, operatorOr:
|
||||
return c.convertLogicalCondition(cond, counter)
|
||||
case operatorIn, operatorNotIn:
|
||||
return c.convertInCondition(cond, counter)
|
||||
case operatorBetween:
|
||||
return c.convertBetweenCondition(cond, counter)
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported operator: %v", cond.Operator)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *filter) convertInCondition(
|
||||
cond *universalFilterCondition,
|
||||
counter *int,
|
||||
) (*convertResult, error) {
|
||||
condField := cond.Field
|
||||
if condField == "" || cond.Value == nil {
|
||||
return nil, fmt.Errorf("milvus filter condition is nil")
|
||||
}
|
||||
|
||||
s := reflect.ValueOf(cond.Value)
|
||||
if s.Kind() != reflect.Slice || s.Len() <= 0 {
|
||||
return nil, fmt.Errorf("in operator value must be a slice with at least one value: %v", cond.Value)
|
||||
}
|
||||
|
||||
paramName := c.convertParamName(cond.Field, counter)
|
||||
return &convertResult{
|
||||
exprStr: fmt.Sprintf("%s %s {%s}", condField, strings.ToLower(cond.Operator), paramName),
|
||||
params: map[string]any{paramName: cond.Value},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c *filter) convertBetweenCondition(
|
||||
cond *universalFilterCondition,
|
||||
counter *int,
|
||||
) (*convertResult, error) {
|
||||
condField := cond.Field
|
||||
if condField == "" || cond.Value == nil {
|
||||
return nil, fmt.Errorf("milvus filter condition is nil")
|
||||
}
|
||||
|
||||
value := reflect.ValueOf(cond.Value)
|
||||
if value.Kind() != reflect.Slice || value.Len() != 2 {
|
||||
return nil, fmt.Errorf("between operator value must be a slice with two elements: %v", cond.Value)
|
||||
}
|
||||
|
||||
paramBase := c.convertParamName(cond.Field, counter)
|
||||
paramName1 := fmt.Sprintf("%s_%d", paramBase, 0)
|
||||
paramName2 := fmt.Sprintf("%s_%d", paramBase, 1)
|
||||
return &convertResult{
|
||||
exprStr: fmt.Sprintf("%s >= {%s} and %s <= {%s}", condField, paramName1, condField, paramName2),
|
||||
params: map[string]any{
|
||||
paramName1: value.Index(0).Interface(),
|
||||
paramName2: value.Index(1).Interface(),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func formatValue(value any) string {
|
||||
switch v := value.(type) {
|
||||
case string:
|
||||
return fmt.Sprintf("\"%s\"", escapeDoubleQuotes(v))
|
||||
case int, int8, int16, int32, int64:
|
||||
return fmt.Sprintf("%d", v)
|
||||
case uint, uint8, uint16, uint32, uint64:
|
||||
return fmt.Sprintf("%d", v)
|
||||
case float32, float64:
|
||||
return fmt.Sprintf("%v", v)
|
||||
case bool:
|
||||
if v {
|
||||
return "true"
|
||||
}
|
||||
return "false"
|
||||
case time.Time:
|
||||
return fmt.Sprintf("%d", v.Unix())
|
||||
default:
|
||||
return fmt.Sprintf("\"%v\"", value)
|
||||
}
|
||||
}
|
||||
|
||||
// escapeDoubleQuotes escapes double quotes in a string for use in Milvus expressions.
|
||||
func escapeDoubleQuotes(s string) string {
|
||||
return strings.ReplaceAll(s, "\"", "\\\"")
|
||||
}
|
||||
|
||||
// convertParamName converts field name to a valid Milvus template parameter name.
|
||||
// Milvus template parameters don't support '.' character, so we replace it with '_'.
|
||||
func (c *filter) convertParamName(field string, counter *int) string {
|
||||
*counter++
|
||||
return fmt.Sprintf("%s_%d", strings.ReplaceAll(field, ".", "_"), *counter)
|
||||
}
|
||||
|
||||
type universalFilterCondition struct {
|
||||
Field string `json:"field,omitempty" jsonschema:"description=The metadata field to filter on (required for comparison operators)"`
|
||||
Operator string `json:"operator" jsonschema:"description=The operator to use,enum=eq,enum=ne,enum=gt,enum=gte,enum=lt,enum=lte,enum=in,enum=not in,enum=like,enum=not like,enum=between,enum=and,enum=or"`
|
||||
Value any `json:"value,omitempty" jsonschema:"description=The value to compare against (for comparison operators) or array of sub-conditions (for logical operators and/or)"`
|
||||
}
|
||||
|
||||
func (c *universalFilterCondition) UnmarshalJSON(data []byte) error {
|
||||
type Alias struct {
|
||||
Field string `json:"field,omitempty"`
|
||||
Operator string `json:"operator"`
|
||||
Value any `json:"value,omitempty"`
|
||||
}
|
||||
|
||||
var aux Alias
|
||||
if err := json.Unmarshal(data, &aux); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
c.Field = aux.Field
|
||||
c.Operator = strings.ToLower(aux.Operator)
|
||||
|
||||
// Handle logical operators (and/or) - Value should be []*UniversalFilterCondition
|
||||
if c.Operator == operatorAnd || c.Operator == operatorOr {
|
||||
// Value can be an array of conditions
|
||||
valueSlice, ok := aux.Value.([]any)
|
||||
if !ok {
|
||||
return fmt.Errorf("logical operator %s requires an array of conditions", c.Operator)
|
||||
}
|
||||
|
||||
conditions := make([]*universalFilterCondition, 0, len(valueSlice))
|
||||
for i, v := range valueSlice {
|
||||
condBytes, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal condition at index %d: %w", i, err)
|
||||
}
|
||||
|
||||
var cond universalFilterCondition
|
||||
if err := json.Unmarshal(condBytes, &cond); err != nil {
|
||||
return fmt.Errorf("failed to unmarshal condition at index %d: %w", i, err)
|
||||
}
|
||||
conditions = append(conditions, &cond)
|
||||
}
|
||||
c.Value = conditions
|
||||
} else {
|
||||
c.Value = aux.Value
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// MarshalJSON implements custom JSON marshaling for UniversalFilterCondition.
|
||||
func (c *universalFilterCondition) MarshalJSON() ([]byte, error) {
|
||||
type Alias struct {
|
||||
Field string `json:"field,omitempty"`
|
||||
Operator string `json:"operator"`
|
||||
Value any `json:"value,omitempty"`
|
||||
}
|
||||
|
||||
return json.Marshal(&Alias{
|
||||
Field: c.Field,
|
||||
Operator: c.Operator,
|
||||
Value: c.Value,
|
||||
})
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,243 @@
|
||||
package milvus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/milvus-io/milvus/client/v2/column"
|
||||
"github.com/milvus-io/milvus/client/v2/entity"
|
||||
client "github.com/milvus-io/milvus/client/v2/milvusclient"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestContentFieldUsesChineseAnalyzer(t *testing.T) {
|
||||
field := newContentField().ProtoMessage()
|
||||
require.Equal(t, fieldContent, field.GetName())
|
||||
require.Equal(t, int32(entity.FieldTypeVarChar), int32(field.GetDataType()))
|
||||
params := make(map[string]string)
|
||||
for _, param := range field.GetTypeParams() {
|
||||
params[param.GetKey()] = param.GetValue()
|
||||
}
|
||||
require.Equal(t, "65535", params["max_length"])
|
||||
require.Equal(t, "true", params["enable_analyzer"])
|
||||
require.Equal(t, "true", params["enable_match"])
|
||||
require.JSONEq(t, `{"type":"chinese"}`, params["analyzer_params"])
|
||||
}
|
||||
|
||||
func TestUpdateChunkEnabledStatusInCollectionSkipsEmptyChunkIDs(t *testing.T) {
|
||||
repo := &milvusRepository{}
|
||||
|
||||
require.NoError(t, repo.updateChunkEnabledStatusInCollection(
|
||||
context.Background(),
|
||||
"weknora_embeddings_1024",
|
||||
nil,
|
||||
false,
|
||||
))
|
||||
require.NoError(t, repo.updateChunkEnabledStatusInCollection(
|
||||
context.Background(),
|
||||
"weknora_embeddings_1024",
|
||||
[]string{},
|
||||
true,
|
||||
))
|
||||
}
|
||||
|
||||
func TestUpdateChunkEnabledStatusInCollectionsPropagatesFailure(t *testing.T) {
|
||||
wantErr := errors.New("upsert failed")
|
||||
err := updateChunkEnabledStatusInCollections(
|
||||
context.Background(),
|
||||
[]string{"other_collection", "weknora_embeddings_1024"},
|
||||
"weknora_embeddings",
|
||||
nil,
|
||||
[]string{"chunk-1"},
|
||||
func(_ context.Context, collection string, _ []string, enabled bool) error {
|
||||
if collection == "weknora_embeddings_1024" && !enabled {
|
||||
return wantErr
|
||||
}
|
||||
return nil
|
||||
},
|
||||
)
|
||||
require.ErrorIs(t, err, wantErr)
|
||||
}
|
||||
|
||||
func TestMatchesCollection(t *testing.T) {
|
||||
for _, name := range []string{"weknora_embeddings_1", "weknora_embeddings_1024"} {
|
||||
require.True(t, matchesCollection(name, "weknora_embeddings"), name)
|
||||
}
|
||||
for _, name := range []string{
|
||||
"weknora_embeddings", "weknora_embeddings_", "weknora_embeddings_0",
|
||||
"weknora_embeddings_-1", "weknora_embeddings_+1", "weknora_embeddings_01024",
|
||||
"weknora_embeddings_1024_backup", "weknora_embeddings_backup_1024",
|
||||
"weknora_embeddings_other_tenant_1024", "weknora_embeddings2_1024",
|
||||
"weknora_embeddings_9999999999999999999999999", "other_1024", "_1024",
|
||||
} {
|
||||
require.False(t, matchesCollection(name, "weknora_embeddings"), name)
|
||||
}
|
||||
require.False(t, matchesCollection("_1024", ""))
|
||||
}
|
||||
|
||||
func TestUpdateChunkEnabledStatusDoesNotCrossNamespaces(t *testing.T) {
|
||||
var touched []string
|
||||
err := updateChunkEnabledStatusInCollections(
|
||||
context.Background(),
|
||||
[]string{"weknora_embeddings_1024_backup", "weknora_embeddings_other_1024", "weknora_embeddings_1024"},
|
||||
"weknora_embeddings", nil, []string{"chunk-1"},
|
||||
func(_ context.Context, collection string, _ []string, _ bool) error {
|
||||
touched = append(touched, collection)
|
||||
return nil
|
||||
},
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"weknora_embeddings_1024", "weknora_embeddings_1024"}, touched)
|
||||
}
|
||||
|
||||
func keywordResult(chunkID string, score float32) []client.ResultSet {
|
||||
return []client.ResultSet{{
|
||||
ResultCount: 1,
|
||||
Scores: []float32{score},
|
||||
Fields: client.DataSet{
|
||||
column.NewColumnVarChar(fieldID, []string{"index-" + chunkID}),
|
||||
column.NewColumnVarChar(fieldChunkID, []string{chunkID}),
|
||||
column.NewColumnVarChar(fieldContent, []string{"test content"}),
|
||||
},
|
||||
}}
|
||||
}
|
||||
|
||||
func TestKeywordsRetrieveReturnsAllCollectionFailures(t *testing.T) {
|
||||
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
|
||||
missingSparseField := errors.New("field content_sparse does not exist")
|
||||
conversionFailure := errors.New("cannot decode search result")
|
||||
var searched []string
|
||||
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
|
||||
Query: "test query", TopK: 10,
|
||||
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024", "weknora_embeddings_backup_1024"},
|
||||
func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
|
||||
request, err := option.Request()
|
||||
require.NoError(t, err)
|
||||
searched = append(searched, request.CollectionName)
|
||||
if request.CollectionName == "weknora_embeddings_384" {
|
||||
return nil, missingSparseField
|
||||
}
|
||||
return []client.ResultSet{{Err: conversionFailure}}, nil
|
||||
})
|
||||
require.Nil(t, results)
|
||||
require.ErrorIs(t, err, missingSparseField)
|
||||
require.ErrorIs(t, err, conversionFailure)
|
||||
require.ErrorContains(t, err, "all 2 matching Milvus collections")
|
||||
require.ErrorContains(t, err, "weknora_embeddings_384")
|
||||
require.ErrorContains(t, err, "weknora_embeddings_1024")
|
||||
require.Equal(t, []string{"weknora_embeddings_384", "weknora_embeddings_1024"}, searched)
|
||||
}
|
||||
|
||||
func TestKeywordsRetrievePreservesBM25AndRanksAcrossCollections(t *testing.T) {
|
||||
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
|
||||
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
|
||||
Query: "test query", TopK: 1, KnowledgeBaseIDs: []string{"kb-allowed"},
|
||||
TagIDs: []string{"tag-allowed"}, ExcludeChunkIDs: []string{"chunk-excluded"},
|
||||
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024", "weknora_embeddings_other_1536"},
|
||||
func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
|
||||
request, err := option.Request()
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, request.Dsl, fieldKnowledgeBaseID)
|
||||
require.Contains(t, request.Dsl, fieldTagID)
|
||||
require.Contains(t, request.Dsl, fieldChunkID)
|
||||
require.Contains(t, request.Dsl, fieldIsEnabled)
|
||||
require.NotEmpty(t, request.ExprTemplateValues)
|
||||
for _, param := range request.SearchParams {
|
||||
if param.Key == "anns_field" {
|
||||
require.Equal(t, fieldContentSparse, param.Value)
|
||||
}
|
||||
}
|
||||
switch request.CollectionName {
|
||||
case "weknora_embeddings_384":
|
||||
return keywordResult("low-score", 2.5), nil
|
||||
case "weknora_embeddings_1024":
|
||||
return keywordResult("high-score", 17.25), nil
|
||||
default:
|
||||
t.Fatalf("searched unrelated collection %s", request.CollectionName)
|
||||
return nil, nil
|
||||
}
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
require.Equal(t, types.KeywordsRetrieverType, results[0].RetrieverType)
|
||||
require.Len(t, results[0].Results, 1)
|
||||
require.Equal(t, "high-score", results[0].Results[0].ChunkID)
|
||||
require.Equal(t, 17.25, results[0].Results[0].Score)
|
||||
require.Equal(t, types.MatchTypeKeywords, results[0].Results[0].MatchType)
|
||||
}
|
||||
|
||||
func TestKeywordsRetrieveKeepsSuccessfulCollections(t *testing.T) {
|
||||
for _, empty := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "results", true: "empty"}[empty], func(t *testing.T) {
|
||||
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
|
||||
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
|
||||
Query: "query", TopK: 5,
|
||||
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024"},
|
||||
func(_ context.Context, option client.SearchOption) ([]client.ResultSet, error) {
|
||||
request, err := option.Request()
|
||||
require.NoError(t, err)
|
||||
if request.CollectionName == "weknora_embeddings_384" {
|
||||
return nil, errors.New("missing sparse field")
|
||||
}
|
||||
if empty {
|
||||
return nil, nil
|
||||
}
|
||||
return keywordResult("available", 8), nil
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
if empty {
|
||||
require.Empty(t, results[0].Results)
|
||||
} else {
|
||||
require.Len(t, results[0].Results, 1)
|
||||
require.Equal(t, "available", results[0].Results[0].ChunkID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeywordsRetrieveRejectsMissingScores(t *testing.T) {
|
||||
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
|
||||
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
|
||||
Query: "query", TopK: 5,
|
||||
}, []string{"weknora_embeddings_384"},
|
||||
func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
|
||||
result := keywordResult("missing-score", 0)
|
||||
result[0].Scores = nil
|
||||
return result, nil
|
||||
})
|
||||
require.Nil(t, results)
|
||||
require.ErrorContains(t, err, "expected 1 BM25 scores, received 0")
|
||||
}
|
||||
|
||||
func TestKeywordsRetrieveNoMatchingCollections(t *testing.T) {
|
||||
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
|
||||
results, err := repo.keywordsRetrieveInCollections(context.Background(), types.RetrieveParams{
|
||||
Query: "query", TopK: 5,
|
||||
}, []string{"weknora_embeddings_backup_384"},
|
||||
func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
|
||||
t.Fatal("must not search unrelated collection")
|
||||
return nil, nil
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, results, 1)
|
||||
require.Empty(t, results[0].Results)
|
||||
}
|
||||
|
||||
func TestKeywordsRetrieveHonorsCancellation(t *testing.T) {
|
||||
repo := &milvusRepository{collectionBaseName: "weknora_embeddings"}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
results, err := repo.keywordsRetrieveInCollections(ctx, types.RetrieveParams{
|
||||
Query: "query", TopK: 5,
|
||||
}, []string{"weknora_embeddings_384", "weknora_embeddings_1024"},
|
||||
func(context.Context, client.SearchOption) ([]client.ResultSet, error) {
|
||||
cancel()
|
||||
return keywordResult("cancelled", 8), nil
|
||||
})
|
||||
require.Nil(t, results)
|
||||
require.ErrorIs(t, err, context.Canceled)
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package milvus
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/milvus-io/milvus/client/v2/entity"
|
||||
client "github.com/milvus-io/milvus/client/v2/milvusclient"
|
||||
)
|
||||
|
||||
type milvusRepository struct {
|
||||
filter
|
||||
client *client.Client
|
||||
collectionBaseName string
|
||||
metricType entity.MetricType
|
||||
shardsNum int // 0 = use Milvus default (1)
|
||||
replicaNumber int // 0 = use Milvus default (1); set at LoadCollection time
|
||||
// Cache for initialized collections (dimension -> true)
|
||||
initializedCollections sync.Map
|
||||
}
|
||||
|
||||
type MilvusVectorEmbedding struct {
|
||||
ID string `json:"id"`
|
||||
Content string `json:"content"`
|
||||
SourceID string `json:"source_id"`
|
||||
SourceType int `json:"source_type"`
|
||||
ChunkID string `json:"chunk_id"`
|
||||
KnowledgeID string `json:"knowledge_id"`
|
||||
KnowledgeBaseID string `json:"knowledge_base_id"`
|
||||
TagID string `json:"tag_id"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
IsEnabled bool `json:"is_enabled"`
|
||||
}
|
||||
|
||||
type MilvusVectorEmbeddingWithScore struct {
|
||||
MilvusVectorEmbedding
|
||||
Score float64
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package milvus
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/milvus-io/milvus/client/v2/column"
|
||||
client "github.com/milvus-io/milvus/client/v2/milvusclient"
|
||||
)
|
||||
|
||||
func TestCopyResultPreservesNativeFloatVector(t *testing.T) {
|
||||
want := []float32{0.1, 0.25, -0.5}
|
||||
for _, field := range []column.Column{
|
||||
column.NewColumnFloatVector(fieldEmbedding, 3, [][]float32{want}),
|
||||
column.NewColumnDoubleArray(fieldEmbedding, [][]float64{{0.1, 0.25, -0.5}}),
|
||||
} {
|
||||
result, _, err := convertResultSet([]client.ResultSet{{ResultCount: 1, Fields: client.DataSet{column.NewColumnVarChar(fieldID, []string{"source"}), field}}})
|
||||
if err != nil || len(result) != 1 || !reflect.DeepEqual(result[0].Embedding, want) {
|
||||
t.Fatalf("copy dropped vector from %T: %#v, %v", field, result, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopyResultRejectsUnknownVectorColumn(t *testing.T) {
|
||||
_, _, err := convertResultSet([]client.ResultSet{{ResultCount: 1, Fields: client.DataSet{column.NewColumnVarChar(fieldEmbedding, []string{"invalid"})}}})
|
||||
if err == nil {
|
||||
t.Fatal("invalid vector was silently converted into a zero-dimensional embedding")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
package neo4j
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/neo4j/neo4j-go-driver/v6/neo4j"
|
||||
)
|
||||
|
||||
// Neo4jRepository is a repository for Neo4j
|
||||
type Neo4jRepository struct {
|
||||
driver neo4j.Driver
|
||||
nodePrefix string
|
||||
}
|
||||
|
||||
// NewNeo4jRepository creates a new Neo4j repository
|
||||
func NewNeo4jRepository(driver neo4j.Driver) interfaces.RetrieveGraphRepository {
|
||||
return &Neo4jRepository{driver: driver, nodePrefix: "ENTITY"}
|
||||
}
|
||||
|
||||
// _remove_hyphen removes hyphens from a string
|
||||
func _remove_hyphen(s string) string {
|
||||
return strings.ReplaceAll(s, "-", "_")
|
||||
}
|
||||
|
||||
// Labels returns the labels for a namespace
|
||||
func (n *Neo4jRepository) Labels(namespace types.NameSpace) []string {
|
||||
res := make([]string, 0)
|
||||
for _, label := range namespace.Labels() {
|
||||
res = append(res, n.nodePrefix+_remove_hyphen(label))
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// Label returns the label for a namespace
|
||||
func (n *Neo4jRepository) Label(namespace types.NameSpace) string {
|
||||
labels := n.Labels(namespace)
|
||||
return strings.Join(labels, ":")
|
||||
}
|
||||
|
||||
// AddGraph adds a graph to the Neo4j repository
|
||||
func (n *Neo4jRepository) AddGraph(ctx context.Context, namespace types.NameSpace, graphs []*types.GraphData) error {
|
||||
if n.driver == nil {
|
||||
logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
|
||||
return nil
|
||||
}
|
||||
for _, graph := range graphs {
|
||||
if err := n.addGraph(ctx, namespace, graph); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// addGraph adds a graph to the Neo4j repository
|
||||
func (n *Neo4jRepository) addGraph(ctx context.Context, namespace types.NameSpace, graph *types.GraphData) error {
|
||||
session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
|
||||
defer session.Close(ctx)
|
||||
|
||||
_, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) {
|
||||
// Node import query
|
||||
node_import_query := `
|
||||
UNWIND $data AS row
|
||||
CALL apoc.merge.node(row.labels, {name: row.name, kg: row.knowledge_id}, row.props, {}) YIELD node
|
||||
SET node.chunks = apoc.coll.union(node.chunks, row.chunks)
|
||||
RETURN distinct 'done' AS result
|
||||
`
|
||||
nodeData := []map[string]interface{}{}
|
||||
for _, node := range graph.Node {
|
||||
nodeData = append(nodeData, map[string]interface{}{
|
||||
"name": node.Name,
|
||||
"knowledge_id": namespace.Knowledge,
|
||||
"props": map[string][]string{"attributes": node.Attributes},
|
||||
"chunks": node.Chunks,
|
||||
"labels": n.Labels(namespace),
|
||||
})
|
||||
}
|
||||
if _, err := tx.Run(ctx, node_import_query, map[string]interface{}{"data": nodeData}); err != nil {
|
||||
return nil, fmt.Errorf("failed to create nodes: %v", err)
|
||||
}
|
||||
|
||||
// Relationship import query
|
||||
rel_import_query := `
|
||||
UNWIND $data AS row
|
||||
CALL apoc.merge.node(row.source_labels, {name: row.source, kg: row.knowledge_id}, {}, {}) YIELD node as source
|
||||
CALL apoc.merge.node(row.target_labels, {name: row.target, kg: row.knowledge_id}, {}, {}) YIELD node as target
|
||||
CALL apoc.merge.relationship(source, row.type, {}, row.attributes, target) YIELD rel
|
||||
RETURN distinct 'done'
|
||||
`
|
||||
relData := []map[string]interface{}{}
|
||||
for _, rel := range graph.Relation {
|
||||
relData = append(relData, map[string]interface{}{
|
||||
"source": rel.Node1,
|
||||
"target": rel.Node2,
|
||||
"knowledge_id": namespace.Knowledge,
|
||||
"type": rel.Type,
|
||||
"source_labels": n.Labels(namespace),
|
||||
"target_labels": n.Labels(namespace),
|
||||
})
|
||||
}
|
||||
if _, err := tx.Run(ctx, rel_import_query, map[string]interface{}{"data": relData}); err != nil {
|
||||
return nil, fmt.Errorf("failed to create relationships: %v", err)
|
||||
}
|
||||
return nil, nil
|
||||
})
|
||||
if err != nil {
|
||||
logger.Errorf(ctx, "failed to add graph: %v", err)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DelGraph deletes a graph from the Neo4j repository
|
||||
func (n *Neo4jRepository) DelGraph(ctx context.Context, namespaces []types.NameSpace) error {
|
||||
if n.driver == nil {
|
||||
logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
|
||||
return nil
|
||||
}
|
||||
session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeWrite})
|
||||
defer session.Close(ctx)
|
||||
|
||||
result, err := session.ExecuteWrite(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) {
|
||||
for _, namespace := range namespaces {
|
||||
labelExpr := n.Label(namespace)
|
||||
|
||||
deleteRelsQuery := `
|
||||
CALL apoc.periodic.iterate(
|
||||
"MATCH (n:` + labelExpr + ` {kg: $knowledge_id})-[r]-(m:` + labelExpr + ` {kg: $knowledge_id}) RETURN r",
|
||||
"DELETE r",
|
||||
{batchSize: 1000, parallel: true, params: {knowledge_id: $knowledge_id}}
|
||||
) YIELD batches, total
|
||||
RETURN total
|
||||
`
|
||||
if _, err := tx.Run(ctx, deleteRelsQuery, map[string]interface{}{"knowledge_id": namespace.Knowledge}); err != nil {
|
||||
return nil, fmt.Errorf("failed to delete relationships: %v", err)
|
||||
}
|
||||
|
||||
deleteNodesQuery := `
|
||||
CALL apoc.periodic.iterate(
|
||||
"MATCH (n:` + labelExpr + ` {kg: $knowledge_id}) RETURN n",
|
||||
"DELETE n",
|
||||
{batchSize: 1000, parallel: true, params: {knowledge_id: $knowledge_id}}
|
||||
) YIELD batches, total
|
||||
RETURN total
|
||||
`
|
||||
if _, err := tx.Run(ctx, deleteNodesQuery, map[string]interface{}{"knowledge_id": namespace.Knowledge}); err != nil {
|
||||
return nil, fmt.Errorf("failed to delete nodes: %v", err)
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
logger.Infof(ctx, "delete graph result: %v", result)
|
||||
return nil
|
||||
}
|
||||
|
||||
// SearchNode searches for nodes in the Neo4j repository
|
||||
func (n *Neo4jRepository) SearchNode(
|
||||
ctx context.Context,
|
||||
namespace types.NameSpace,
|
||||
nodes []string,
|
||||
) (*types.GraphData, error) {
|
||||
if n.driver == nil {
|
||||
logger.Warnf(ctx, "NOT SUPPORT RETRIEVE GRAPH")
|
||||
return nil, nil
|
||||
}
|
||||
session := n.driver.NewSession(ctx, neo4j.SessionConfig{AccessMode: neo4j.AccessModeRead})
|
||||
defer session.Close(ctx)
|
||||
|
||||
result, err := session.ExecuteRead(ctx, func(tx neo4j.ManagedTransaction) (interface{}, error) {
|
||||
labelExpr := n.Label(namespace)
|
||||
query := `
|
||||
MATCH (n:` + labelExpr + `)-[r]-(m:` + labelExpr + `)
|
||||
WHERE ANY(nodeText IN $nodes WHERE n.name CONTAINS nodeText)
|
||||
RETURN n, r, m
|
||||
`
|
||||
params := map[string]interface{}{"nodes": nodes}
|
||||
result, err := tx.Run(ctx, query, params)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to run query: %v", err)
|
||||
}
|
||||
|
||||
graphData := &types.GraphData{}
|
||||
nodeSeen := make(map[string]bool)
|
||||
for result.Next(ctx) {
|
||||
record := result.Record()
|
||||
node, _ := record.Get("n")
|
||||
rel, _ := record.Get("r")
|
||||
targetNode, _ := record.Get("m")
|
||||
|
||||
nodeData := node.(neo4j.Node)
|
||||
targetNodeData := targetNode.(neo4j.Node)
|
||||
|
||||
// Convert node to types.Node
|
||||
for _, n := range []neo4j.Node{nodeData, targetNodeData} {
|
||||
nameStr := n.Props["name"].(string)
|
||||
if _, ok := nodeSeen[nameStr]; !ok {
|
||||
nodeSeen[nameStr] = true
|
||||
graphData.Node = append(graphData.Node, &types.GraphNode{
|
||||
Name: nameStr,
|
||||
Chunks: listI2listS(n.Props["chunks"].([]interface{})),
|
||||
Attributes: listI2listS(n.Props["attributes"].([]interface{})),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Convert relationship to types.Relation
|
||||
relData := rel.(neo4j.Relationship)
|
||||
graphData.Relation = append(graphData.Relation, &types.GraphRelation{
|
||||
Node1: nodeData.Props["name"].(string),
|
||||
Node2: targetNodeData.Props["name"].(string),
|
||||
Type: relData.Type,
|
||||
})
|
||||
}
|
||||
return graphData, nil
|
||||
})
|
||||
if err != nil {
|
||||
logger.Errorf(ctx, "search node failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
return result.(*types.GraphData), nil
|
||||
}
|
||||
|
||||
func listI2listS(list []any) []string {
|
||||
result := make([]string, len(list))
|
||||
for i, v := range list {
|
||||
result[i] = fmt.Sprintf("%v", v)
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package opensearch
|
||||
|
||||
import "context"
|
||||
|
||||
// AuditSink receives audit events emitted from within the driver at the
|
||||
// exact moment they occur (index provisioned, reindex executed). The driver
|
||||
// owns this abstraction so it imports no service package — the dependency
|
||||
// arrow stays one-way (service implements AuditSink; the driver only invokes
|
||||
// it). A nil sink is a no-op, so tests and the env-path (no audit service)
|
||||
// need no special casing.
|
||||
type AuditSink interface {
|
||||
// EmitIndexCreated fires once when the driver provisions a new k-NN
|
||||
// index. alias is the per-dimension alias (or the keyword-only index
|
||||
// name); dim is the embedding dimension (0 for the keyword-only index).
|
||||
EmitIndexCreated(ctx context.Context, alias string, dim int)
|
||||
// EmitReindexExecuted fires when CopyIndices finishes copying docs from
|
||||
// one index to another.
|
||||
EmitReindexExecuted(ctx context.Context, srcAlias, dstAlias string, docs int64)
|
||||
}
|
||||
|
||||
// nopSink is the null-object used when no sink is configured, so emit call
|
||||
// sites never need a nil check.
|
||||
type nopSink struct{}
|
||||
|
||||
func (nopSink) EmitIndexCreated(context.Context, string, int) {}
|
||||
func (nopSink) EmitReindexExecuted(context.Context, string, string, int64) {}
|
||||
|
||||
var _ AuditSink = nopSink{}
|
||||
|
||||
// Option configures a Repository at construction time.
|
||||
type Option func(*Repository)
|
||||
|
||||
// WithAuditSink injects an audit sink. A nil sink is ignored (the Repository
|
||||
// keeps its default no-op behavior).
|
||||
func WithAuditSink(s AuditSink) Option {
|
||||
return func(r *Repository) {
|
||||
if s != nil {
|
||||
r.sink = s
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// auditSink returns the configured sink, or a no-op if none was set (e.g. a
|
||||
// Repository built directly in tests, or constructed without WithAuditSink).
|
||||
func (r *Repository) auditSink() AuditSink {
|
||||
if r.sink == nil {
|
||||
return nopSink{}
|
||||
}
|
||||
return r.sink
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// spySink records audit events for assertions.
|
||||
type spySink struct {
|
||||
mu sync.Mutex
|
||||
indexCreated []indexCreatedEvent
|
||||
reindex []reindexEvent
|
||||
}
|
||||
|
||||
type indexCreatedEvent struct {
|
||||
alias string
|
||||
dim int
|
||||
}
|
||||
type reindexEvent struct {
|
||||
src, dst string
|
||||
docs int64
|
||||
}
|
||||
|
||||
func (s *spySink) EmitIndexCreated(_ context.Context, alias string, dim int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.indexCreated = append(s.indexCreated, indexCreatedEvent{alias, dim})
|
||||
}
|
||||
|
||||
func (s *spySink) EmitReindexExecuted(_ context.Context, src, dst string, docs int64) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.reindex = append(s.reindex, reindexEvent{src, dst, docs})
|
||||
}
|
||||
|
||||
// TestBuildInternalCfg_ReadsHNSWFields verifies the HNSW IndexConfig fields
|
||||
// flow into internalCfg, falling back to defaults when unset.
|
||||
func TestBuildInternalCfg_ReadsHNSWFields(t *testing.T) {
|
||||
t.Run("reads set fields", func(t *testing.T) {
|
||||
cfg, err := buildInternalCfg(&types.IndexConfig{
|
||||
HNSWM: 24,
|
||||
HNSWEFConstruction: 200,
|
||||
HNSWEFSearch: 128,
|
||||
KNNEngine: "faiss",
|
||||
NumberOfShards: 2,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.hnswM != 24 || cfg.hnswEFConstruction != 200 || cfg.efSearch != 128 || cfg.knnEngine != "faiss" {
|
||||
t.Errorf("HNSW not wired: %+v", cfg)
|
||||
}
|
||||
if cfg.shards != 2 {
|
||||
t.Errorf("shards: want 2, got %d", cfg.shards)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("defaults when unset", func(t *testing.T) {
|
||||
cfg, err := buildInternalCfg(&types.IndexConfig{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.hnswM != 16 || cfg.hnswEFConstruction != 100 || cfg.efSearch != 100 || cfg.knnEngine != "lucene" {
|
||||
t.Errorf("defaults not applied: %+v", cfg)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("nil config is all defaults", func(t *testing.T) {
|
||||
cfg, err := buildInternalCfg(nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if cfg.knnEngine != "lucene" || cfg.hnswM != 16 {
|
||||
t.Errorf("nil defaults wrong: %+v", cfg)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestBuildIndexMapping_ReflectsHNSWConfig verifies the end-to-end path
|
||||
// IndexConfig → buildInternalCfg → buildIndexMapping carries the operator's
|
||||
// HNSW values into the cluster mapping JSON (regression guard for the wire-
|
||||
// through, since defaults coincide with common values).
|
||||
func TestBuildIndexMapping_ReflectsHNSWConfig(t *testing.T) {
|
||||
cfg, err := buildInternalCfg(&types.IndexConfig{
|
||||
HNSWM: 24,
|
||||
HNSWEFConstruction: 200,
|
||||
HNSWEFSearch: 128,
|
||||
KNNEngine: "faiss",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
body, err := buildIndexMapping(cfg, 768)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s := string(body)
|
||||
for _, want := range []string{`"m":24`, `"ef_construction":200`, `"engine":"faiss"`, `"knn.algo_param.ef_search":128`} {
|
||||
if !strings.Contains(s, want) {
|
||||
t.Errorf("mapping JSON missing %q\n%s", want, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestWithAuditSink_SetsAndNilSafe verifies the functional option.
|
||||
func TestWithAuditSink_SetsAndNilSafe(t *testing.T) {
|
||||
spy := &spySink{}
|
||||
r := &Repository{}
|
||||
WithAuditSink(spy)(r)
|
||||
if r.sink != spy {
|
||||
t.Fatal("WithAuditSink did not set the sink")
|
||||
}
|
||||
// nil must not clobber an already-set sink
|
||||
WithAuditSink(nil)(r)
|
||||
if r.sink != spy {
|
||||
t.Fatal("WithAuditSink(nil) clobbered the sink")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditSink_NopByDefault verifies a Repository with no sink does not panic
|
||||
// when the audit accessor is used (nopSink fallback).
|
||||
func TestAuditSink_NopByDefault(t *testing.T) {
|
||||
var _ AuditSink = nopSink{} // compile-time assertion
|
||||
r := &Repository{} // sink left nil
|
||||
// Must not panic.
|
||||
r.auditSink().EmitIndexCreated(context.Background(), "weknora_768", 768)
|
||||
r.auditSink().EmitReindexExecuted(context.Background(), "a", "b", 3)
|
||||
}
|
||||
|
||||
// TestAuditSink_EmitIndexCreated_OnEnsureReady verifies createIndexAndAlias
|
||||
// emits exactly one index-created event with the per-dim alias when a new
|
||||
// index is provisioned.
|
||||
func TestAuditSink_EmitIndexCreated_OnEnsureReady(t *testing.T) {
|
||||
repo, ts := newTestRepo(t, (&indexLifecycleHandler{}).ServeHTTP)
|
||||
defer ts.Close()
|
||||
spy := &spySink{}
|
||||
repo.sink = spy
|
||||
|
||||
if err := repo.ensureReady(context.Background(), 768); err != nil {
|
||||
t.Fatalf("ensureReady: %v", err)
|
||||
}
|
||||
if len(spy.indexCreated) != 1 {
|
||||
t.Fatalf("want 1 index_created event, got %d", len(spy.indexCreated))
|
||||
}
|
||||
got := spy.indexCreated[0]
|
||||
if got.alias != "weknora_test_768" || got.dim != 768 {
|
||||
t.Errorf("event mismatch: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"sort"
|
||||
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
)
|
||||
|
||||
// BatchUpdateChunkEnabledStatus flips is_enabled for the given chunks. The
|
||||
// status map is grouped by value into one _update_by_query per distinct value
|
||||
// (mirrors the Qdrant grouping pattern), so the request body carries the
|
||||
// chunk ids via a terms filter and the new value via bound script params —
|
||||
// never per-chunk string interpolation.
|
||||
//
|
||||
// Targets the cross-dim <base>_* pattern: a chunk's embedding dimension is
|
||||
// not known here, and the same chunk_id is unique across the store's dim
|
||||
// indices + the keyword-only index.
|
||||
func (r *Repository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
|
||||
if len(chunkStatusMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
groups := map[bool][]string{}
|
||||
for id, v := range chunkStatusMap {
|
||||
groups[v] = append(groups[v], id)
|
||||
}
|
||||
// Deterministic order (false then true) for predictable behavior/tests.
|
||||
for _, v := range []bool{false, true} {
|
||||
ids := groups[v]
|
||||
if len(ids) == 0 {
|
||||
continue
|
||||
}
|
||||
sort.Strings(ids)
|
||||
if err := r.updateByQueryScript(ctx, ids,
|
||||
"ctx._source.is_enabled = params.v", map[string]any{"v": v}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchUpdateChunkTagID sets tag_id for the given chunks, grouped by tag.
|
||||
func (r *Repository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
|
||||
if len(chunkTagMap) == 0 {
|
||||
return nil
|
||||
}
|
||||
groups := map[string][]string{}
|
||||
for id, tag := range chunkTagMap {
|
||||
groups[tag] = append(groups[tag], id)
|
||||
}
|
||||
tags := make([]string, 0, len(groups))
|
||||
for tag := range groups {
|
||||
tags = append(tags, tag)
|
||||
}
|
||||
sort.Strings(tags)
|
||||
for _, tag := range tags {
|
||||
ids := groups[tag]
|
||||
sort.Strings(ids)
|
||||
if err := r.updateByQueryScript(ctx, ids,
|
||||
"ctx._source.tag_id = params.v", map[string]any{"v": tag}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// updateByQueryScript runs an _update_by_query over the cross-dim <base>_*
|
||||
// pattern, matching the given chunk ids via a terms filter and applying a
|
||||
// constant Painless source with caller values flowing only through bound
|
||||
// params (Painless-injection-safe).
|
||||
func (r *Repository) updateByQueryScript(
|
||||
ctx context.Context, chunkIDs []string, source string, params map[string]any,
|
||||
) error {
|
||||
body, err := json.Marshal(map[string]any{
|
||||
"query": map[string]any{
|
||||
"terms": map[string]any{"chunk_id": chunkIDs},
|
||||
},
|
||||
"script": map[string]any{
|
||||
"lang": "painless",
|
||||
"source": source,
|
||||
"params": params,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("opensearch: marshal update_by_query body: %w", err)
|
||||
}
|
||||
// Q2: UpdateByQueryParams.Refresh is *bool — the wire value "wait_for" is
|
||||
// not expressible via the typed SDK, so we force an immediate refresh.
|
||||
refresh := true
|
||||
resp, err := r.client.UpdateByQuery(ctx, osapi.UpdateByQueryReq{
|
||||
Indices: []string{r.baseIndex + "_*"},
|
||||
Body: bytes.NewReader(body),
|
||||
Params: osapi.UpdateByQueryParams{Refresh: &refresh},
|
||||
})
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp == nil {
|
||||
return nil
|
||||
}
|
||||
defer drainAndClose(resp.Inspect().Response.Body)
|
||||
return inspectByQueryResponse(io.LimitReader(resp.Inspect().Response.Body, 16<<20))
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// inspectByQueryResponse parses an _update_by_query / _delete_by_query
|
||||
// response and surfaces partial failures without leaking cluster-side reason
|
||||
// strings (which may embed document content). Mirrors inspectBulkResponse:
|
||||
// the full reason goes to the debug log only; the returned error carries the
|
||||
// bounded id + type. A non-zero version_conflicts count with no hard failures
|
||||
// is logged as a warning but not treated as an error.
|
||||
func inspectByQueryResponse(body io.Reader) error {
|
||||
var r struct {
|
||||
VersionConflicts int `json:"version_conflicts"`
|
||||
Failures []struct {
|
||||
ID string `json:"id"`
|
||||
Cause struct {
|
||||
Type string `json:"type"`
|
||||
Reason string `json:"reason"`
|
||||
} `json:"cause"`
|
||||
} `json:"failures"`
|
||||
}
|
||||
if err := json.NewDecoder(body).Decode(&r); err != nil {
|
||||
return fmt.Errorf("opensearch: parse by-query response: %w", ErrTransport)
|
||||
}
|
||||
log := logger.GetLogger(context.Background())
|
||||
if len(r.Failures) == 0 {
|
||||
if r.VersionConflicts > 0 {
|
||||
log.Warnf("[OpenSearch] by-query had %d version conflicts (proceeded)", r.VersionConflicts)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
var msgs []string
|
||||
for _, f := range r.Failures {
|
||||
// Full reason → debug log only (may contain document content).
|
||||
log.Debugf("[OpenSearch] by-query failure: id=%s type=%s reason=%s",
|
||||
f.ID, f.Cause.Type, f.Cause.Reason)
|
||||
if len(msgs) < 5 {
|
||||
msgs = append(msgs, fmt.Sprintf("[%s] %s", f.ID, f.Cause.Type))
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("opensearch: by-query partial failure (%d failed, first 5: %s): %w",
|
||||
len(r.Failures), strings.Join(msgs, "; "), ErrTransport)
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package opensearch
|
||||
|
||||
import "github.com/Tencent/WeKnora/internal/types"
|
||||
|
||||
// internalCfg is the driver-internal, immutable view of IndexConfig.
|
||||
// Defaults are chosen so the env-path (no IndexConfig) and DB-store path
|
||||
// (with IndexConfig) produce identical mappings.
|
||||
type internalCfg struct {
|
||||
shards int
|
||||
replicas int
|
||||
knnEngine string // "lucene" | "faiss"
|
||||
hnswM int
|
||||
hnswEFConstruction int
|
||||
efSearch int
|
||||
}
|
||||
|
||||
// buildInternalCfg projects IndexConfig to the driver-internal view,
|
||||
// substituting defaults for unset (zero / empty) fields. Validation of value
|
||||
// ranges (e.g. hnsw_m / ef_construction caps) is a service-layer concern
|
||||
// handled elsewhere (validateOpenSearchIndexConfig at CreateStore); this
|
||||
// function applies defaults only and never rejects. The env-path bypasses
|
||||
// service validation entirely, so the defaults below are its safety net.
|
||||
//
|
||||
// The OpenSearch-specific HNSW fields (knn_engine, hnsw_m,
|
||||
// hnsw_ef_construction, hnsw_ef_search) are read here. They are omitempty on
|
||||
// IndexConfig, so they do not affect other drivers' serialized config.
|
||||
func buildInternalCfg(c *types.IndexConfig) (internalCfg, error) {
|
||||
cfg := internalCfg{
|
||||
shards: 4, // matches the keyword-index default upstream
|
||||
replicas: 1, // assumes >= 2-node cluster
|
||||
knnEngine: "lucene", // OS default; Faiss preferred only at >= 10M docs
|
||||
hnswM: 16, // OS official default
|
||||
hnswEFConstruction: 100, // OS official default
|
||||
efSearch: 100, // OS default
|
||||
}
|
||||
if c == nil {
|
||||
return cfg, nil
|
||||
}
|
||||
if c.NumberOfShards > 0 {
|
||||
cfg.shards = c.NumberOfShards
|
||||
}
|
||||
if c.NumberOfReplicas > 0 {
|
||||
cfg.replicas = c.NumberOfReplicas
|
||||
}
|
||||
if c.KNNEngine != "" {
|
||||
cfg.knnEngine = c.KNNEngine
|
||||
}
|
||||
if c.HNSWM > 0 {
|
||||
cfg.hnswM = c.HNSWM
|
||||
}
|
||||
if c.HNSWEFConstruction > 0 {
|
||||
cfg.hnswEFConstruction = c.HNSWEFConstruction
|
||||
}
|
||||
if c.HNSWEFSearch > 0 {
|
||||
cfg.efSearch = c.HNSWEFSearch
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
@@ -0,0 +1,194 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// copyBatchSize is the pagination size for the source scan. Kept under the
|
||||
// BatchSave per-call cap so each copied page is a single bulk request.
|
||||
const copyBatchSize = 500
|
||||
|
||||
// copySourceDoc is the full _source read during CopyIndices — it includes the
|
||||
// embedding vector and is_recommended, which the retrieve-path hit struct
|
||||
// omits because retrieval does not need them.
|
||||
type copySourceDoc struct {
|
||||
Content string `json:"content"`
|
||||
SourceID string `json:"source_id"`
|
||||
SourceType int `json:"source_type"`
|
||||
ChunkID string `json:"chunk_id"`
|
||||
KnowledgeID string `json:"knowledge_id"`
|
||||
KnowledgeBaseID string `json:"knowledge_base_id"`
|
||||
TagID string `json:"tag_id"`
|
||||
IsEnabled bool `json:"is_enabled"`
|
||||
IsRecommended bool `json:"is_recommended"`
|
||||
Embedding []float32 `json:"embedding"`
|
||||
}
|
||||
|
||||
// transformSourceID mirrors the sibling drivers' source_id remap:
|
||||
// - regular chunk (source_id == chunk_id) → target chunk id
|
||||
// - generated question (source_id == "<chunk>-<q>") → "<targetChunk>-<q>"
|
||||
// - anything else → a fresh uuid
|
||||
func transformSourceID(sourceID, chunkID, targetChunkID string) string {
|
||||
switch {
|
||||
case sourceID == chunkID:
|
||||
return targetChunkID
|
||||
case strings.HasPrefix(sourceID, chunkID+"-"):
|
||||
return targetChunkID + "-" + strings.TrimPrefix(sourceID, chunkID+"-")
|
||||
default:
|
||||
return uuid.New().String()
|
||||
}
|
||||
}
|
||||
|
||||
// CopyIndices copies all docs of one knowledge base into another (within the
|
||||
// same store) by scanning the source and re-saving via BatchSave — mirroring
|
||||
// the Elasticsearch / Qdrant drivers (search→BatchSave), which yields the
|
||||
// source_id transformation and dim/keyword routing for free. Runs
|
||||
// synchronously and paginates; the large-batch background-task path is a
|
||||
// later change.
|
||||
//
|
||||
// NOTE: from/size pagination is bounded by the index's max_result_window
|
||||
// (default 10000). Copies larger than that require the scroll-based async
|
||||
// path (a later change).
|
||||
func (r *Repository) CopyIndices(
|
||||
ctx context.Context,
|
||||
sourceKnowledgeBaseID string,
|
||||
sourceToTargetKBIDMap map[string]string, // keyed by source knowledge_id (mirrors sibling drivers)
|
||||
sourceToTargetChunkIDMap map[string]string,
|
||||
targetKnowledgeBaseID string,
|
||||
dimension int,
|
||||
knowledgeType string,
|
||||
) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
if len(sourceToTargetChunkIDMap) == 0 {
|
||||
log.Warn("[OpenSearch] CopyIndices: empty chunk mapping, skipping")
|
||||
return nil
|
||||
}
|
||||
if dimension <= 0 {
|
||||
return fmt.Errorf("opensearch: CopyIndices requires dim > 0, got %d: %w",
|
||||
dimension, ErrDimensionMismatch)
|
||||
}
|
||||
if err := r.ensureReady(ctx, dimension); err != nil {
|
||||
return err
|
||||
}
|
||||
alias := r.indexAlias(dimension)
|
||||
|
||||
var total int64
|
||||
for from := 0; ; from += copyBatchSize {
|
||||
docs, err := r.copyScanBatch(ctx, alias, sourceKnowledgeBaseID, from, copyBatchSize)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(docs) == 0 {
|
||||
break
|
||||
}
|
||||
infos := make([]*types.IndexInfo, 0, len(docs))
|
||||
embMap := make(map[string][]float32, len(docs))
|
||||
enabledMap := make(map[string]bool, len(docs))
|
||||
for i := range docs {
|
||||
d := &docs[i]
|
||||
targetChunkID, ok := sourceToTargetChunkIDMap[d.ChunkID]
|
||||
if !ok {
|
||||
log.Warnf("[OpenSearch] CopyIndices: source chunk %s not mapped, skipping", d.ChunkID)
|
||||
continue
|
||||
}
|
||||
targetKnowledgeID, ok := sourceToTargetKBIDMap[d.KnowledgeID]
|
||||
if !ok {
|
||||
log.Warnf("[OpenSearch] CopyIndices: source knowledge %s not mapped, skipping", d.KnowledgeID)
|
||||
continue
|
||||
}
|
||||
targetSourceID := transformSourceID(d.SourceID, d.ChunkID, targetChunkID)
|
||||
if len(d.Embedding) > 0 {
|
||||
// BatchSave looks up embeddings by SourceID (lookupEmbedding),
|
||||
// so key by the target source id — not the chunk id, which is
|
||||
// the Elasticsearch driver's convention.
|
||||
embMap[targetSourceID] = d.Embedding
|
||||
}
|
||||
enabledMap[targetChunkID] = d.IsEnabled
|
||||
infos = append(infos, &types.IndexInfo{
|
||||
Content: d.Content,
|
||||
SourceID: targetSourceID,
|
||||
SourceType: types.SourceType(d.SourceType),
|
||||
ChunkID: targetChunkID,
|
||||
KnowledgeID: targetKnowledgeID,
|
||||
KnowledgeBaseID: targetKnowledgeBaseID,
|
||||
KnowledgeType: knowledgeType,
|
||||
TagID: d.TagID,
|
||||
IsEnabled: d.IsEnabled,
|
||||
IsRecommended: d.IsRecommended,
|
||||
})
|
||||
}
|
||||
if len(infos) > 0 {
|
||||
params := map[string]any{
|
||||
"embedding": embMap,
|
||||
"chunk_enabled": enabledMap,
|
||||
}
|
||||
if err := r.BatchSave(ctx, infos, params); err != nil {
|
||||
return fmt.Errorf("opensearch: CopyIndices batch save: %w", err)
|
||||
}
|
||||
total += int64(len(infos))
|
||||
}
|
||||
if len(docs) < copyBatchSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
log.Infof("[OpenSearch] CopyIndices: copied %d docs (KB %s → %s, dim=%d)",
|
||||
total, sourceKnowledgeBaseID, targetKnowledgeBaseID, dimension)
|
||||
r.auditSink().EmitReindexExecuted(ctx, alias, alias, total)
|
||||
return nil
|
||||
}
|
||||
|
||||
// copyScanBatch reads one page of docs belonging to sourceKB from the per-dim
|
||||
// index, decoding the full _source (including the embedding vector).
|
||||
func (r *Repository) copyScanBatch(
|
||||
ctx context.Context, index, sourceKB string, from, size int,
|
||||
) ([]copySourceDoc, error) {
|
||||
body, err := json.Marshal(map[string]any{
|
||||
"from": from,
|
||||
"size": size,
|
||||
"query": map[string]any{
|
||||
"bool": map[string]any{
|
||||
"filter": []any{
|
||||
map[string]any{"term": map[string]any{"knowledge_base_id": sourceKB}},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opensearch: marshal copy scan body: %w", err)
|
||||
}
|
||||
req := osapi.SearchReq{Indices: []string{index}, Body: bytes.NewReader(body)}
|
||||
resp, err := r.client.Search(ctx, &req)
|
||||
if err != nil {
|
||||
if isNotFound(err) {
|
||||
return nil, fmt.Errorf("opensearch: index %s missing: %w", index, ErrIndexNotFound)
|
||||
}
|
||||
return nil, wrapTransport(err)
|
||||
}
|
||||
defer drainAndClose(resp.Inspect().Response.Body)
|
||||
var parsed struct {
|
||||
Hits struct {
|
||||
Hits []struct {
|
||||
Source copySourceDoc `json:"_source"`
|
||||
} `json:"hits"`
|
||||
} `json:"hits"`
|
||||
}
|
||||
if err := json.NewDecoder(io.LimitReader(resp.Inspect().Response.Body, 64<<20)).Decode(&parsed); err != nil {
|
||||
return nil, fmt.Errorf("opensearch: parse copy scan response: %w", ErrTransport)
|
||||
}
|
||||
out := make([]copySourceDoc, len(parsed.Hits.Hits))
|
||||
for i, h := range parsed.Hits.Hits {
|
||||
out[i] = h.Source
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestTransformSourceID covers the generated-question / regular / fallback
|
||||
// branches that CopyIndices uses to remap source_id.
|
||||
func TestTransformSourceID(t *testing.T) {
|
||||
t.Run("regular chunk uses target chunk id", func(t *testing.T) {
|
||||
if got := transformSourceID("chunk1", "chunk1", "tgt1"); got != "tgt1" {
|
||||
t.Errorf("want tgt1, got %s", got)
|
||||
}
|
||||
})
|
||||
t.Run("generated question preserves question id", func(t *testing.T) {
|
||||
if got := transformSourceID("chunk1-q7", "chunk1", "tgt1"); got != "tgt1-q7" {
|
||||
t.Errorf("want tgt1-q7, got %s", got)
|
||||
}
|
||||
})
|
||||
t.Run("unrelated source id gets fresh uuid", func(t *testing.T) {
|
||||
got := transformSourceID("totally-different", "chunk1", "tgt1")
|
||||
if got == "totally-different" || got == "tgt1" || len(got) != 36 {
|
||||
t.Errorf("want fresh uuid, got %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestCopyIndices_EmptyMapping_NoOp verifies an empty chunk map short-circuits
|
||||
// before any HTTP call.
|
||||
func TestCopyIndices_EmptyMapping_NoOp(t *testing.T) {
|
||||
repo, ts := newTestRepo(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Errorf("unexpected HTTP call: %s %s", r.Method, r.URL.Path)
|
||||
})
|
||||
defer ts.Close()
|
||||
err := repo.CopyIndices(context.Background(), "kbSrc", map[string]string{}, map[string]string{}, "kbDst", 768, "manual")
|
||||
if err != nil {
|
||||
t.Fatalf("want nil, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCopyIndices_ScanThenBatchSave verifies the search→BatchSave path:
|
||||
// remaps IDs, keys the embedding by the *target source id* (OpenSearch
|
||||
// BatchSave's lookup key), and emits one reindex audit event.
|
||||
func TestCopyIndices_ScanThenBatchSave(t *testing.T) {
|
||||
var (
|
||||
mu sync.Mutex
|
||||
bulkBody string
|
||||
searchCnt int
|
||||
)
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.Method == http.MethodHead:
|
||||
w.WriteHeader(http.StatusOK) // alias exists → ensureReady no-op
|
||||
case strings.Contains(r.URL.Path, "_search"):
|
||||
mu.Lock()
|
||||
searchCnt++
|
||||
first := searchCnt == 1
|
||||
mu.Unlock()
|
||||
if first {
|
||||
_, _ = w.Write([]byte(`{"hits":{"hits":[
|
||||
{"_source":{"content":"c","source_id":"srcChunk","source_type":1,"chunk_id":"srcChunk","knowledge_id":"srcKnow","knowledge_base_id":"kbSrc","tag_id":"t","is_enabled":true,"is_recommended":false,"embedding":[0.1,0.2,0.3]}}
|
||||
]}}`))
|
||||
} else {
|
||||
_, _ = w.Write([]byte(`{"hits":{"hits":[]}}`))
|
||||
}
|
||||
case strings.HasSuffix(r.URL.Path, "/_bulk"):
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
mu.Lock()
|
||||
bulkBody = string(b)
|
||||
mu.Unlock()
|
||||
_, _ = w.Write([]byte(`{"errors":false,"items":[]}`))
|
||||
default:
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
}
|
||||
}
|
||||
repo, ts := newTestRepo(t, handler)
|
||||
defer ts.Close()
|
||||
spy := &spySink{}
|
||||
repo.sink = spy
|
||||
|
||||
err := repo.CopyIndices(context.Background(), "kbSrc",
|
||||
map[string]string{"srcKnow": "tgtKnow"}, // knowledge_id remap (sourceToTargetKBIDMap is keyed by knowledge_id, mirroring ES)
|
||||
map[string]string{"srcChunk": "tgtChunk"},
|
||||
"kbDst", 768, "manual")
|
||||
if err != nil {
|
||||
t.Fatalf("CopyIndices: %v", err)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if bulkBody == "" {
|
||||
t.Fatal("no bulk request captured")
|
||||
}
|
||||
// Target IDs present, source IDs gone from the written doc.
|
||||
for _, want := range []string{`"chunk_id":"tgtChunk"`, `"knowledge_id":"tgtKnow"`, `"knowledge_base_id":"kbDst"`, `"source_id":"tgtChunk"`} {
|
||||
if !strings.Contains(bulkBody, want) {
|
||||
t.Errorf("bulk body missing %q\n%s", want, bulkBody)
|
||||
}
|
||||
}
|
||||
if strings.Contains(bulkBody, `"knowledge_base_id":"kbSrc"`) {
|
||||
t.Errorf("bulk body leaked source KB id\n%s", bulkBody)
|
||||
}
|
||||
// Embedding written (keyed by target source id internally).
|
||||
if !strings.Contains(bulkBody, "0.1") {
|
||||
t.Errorf("embedding not written\n%s", bulkBody)
|
||||
}
|
||||
if len(spy.reindex) != 1 || spy.reindex[0].docs != 1 {
|
||||
t.Errorf("want 1 reindex event with docs=1, got %+v", spy.reindex)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBatchUpdateChunkEnabledStatus_GroupedUpdateByQuery verifies the status
|
||||
// map is grouped by value into one _update_by_query per distinct value, each
|
||||
// passing chunk ids via terms + the new value via bound script params.
|
||||
func TestBatchUpdateChunkEnabledStatus_GroupedUpdateByQuery(t *testing.T) {
|
||||
var (
|
||||
mu sync.Mutex
|
||||
bodies []string
|
||||
)
|
||||
handler := func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "_update_by_query") {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
mu.Lock()
|
||||
bodies = append(bodies, string(b))
|
||||
mu.Unlock()
|
||||
_, _ = w.Write([]byte(`{"updated":1,"version_conflicts":0,"failures":[]}`))
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
}
|
||||
repo, ts := newTestRepo(t, handler)
|
||||
defer ts.Close()
|
||||
|
||||
err := repo.BatchUpdateChunkEnabledStatus(context.Background(), map[string]bool{
|
||||
"c1": true, "c2": false, "c3": true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("BatchUpdateChunkEnabledStatus: %v", err)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(bodies) != 2 {
|
||||
t.Fatalf("want 2 grouped update_by_query calls (true/false), got %d: %v", len(bodies), bodies)
|
||||
}
|
||||
joined := strings.Join(bodies, "\n")
|
||||
for _, want := range []string{"c1", "c2", "c3", "is_enabled", "params"} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Errorf("update_by_query bodies missing %q\n%s", want, joined)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchUpdateChunkEnabledStatus_Empty_NoOp(t *testing.T) {
|
||||
repo, ts := newTestRepo(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Errorf("unexpected HTTP call: %s %s", r.Method, r.URL.Path)
|
||||
})
|
||||
defer ts.Close()
|
||||
if err := repo.BatchUpdateChunkEnabledStatus(context.Background(), map[string]bool{}); err != nil {
|
||||
t.Fatalf("want nil, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestInspectByQueryResponse covers the success path and the failure path,
|
||||
// asserting cluster-side reason text is NOT surfaced in the returned error.
|
||||
func TestInspectByQueryResponse(t *testing.T) {
|
||||
t.Run("clean response", func(t *testing.T) {
|
||||
body := strings.NewReader(`{"updated":5,"version_conflicts":0,"failures":[]}`)
|
||||
if err := inspectByQueryResponse(body); err != nil {
|
||||
t.Fatalf("want nil, got %v", err)
|
||||
}
|
||||
})
|
||||
t.Run("failures do not leak reason", func(t *testing.T) {
|
||||
body := strings.NewReader(`{"updated":1,"version_conflicts":0,"failures":[
|
||||
{"id":"c9","cause":{"type":"version_conflict_engine_exception","reason":"SECRET document body leaked here"}}
|
||||
]}`)
|
||||
err := inspectByQueryResponse(body)
|
||||
if err == nil {
|
||||
t.Fatal("want error for failures, got nil")
|
||||
}
|
||||
if strings.Contains(err.Error(), "SECRET") {
|
||||
t.Errorf("error leaked cluster reason: %v", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "version_conflict_engine_exception") {
|
||||
t.Errorf("error should surface bounded type: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// additionalParams contract (verified against
|
||||
// elasticsearch/structs.go:ToDBVectorEmbedding and
|
||||
// keywords_vector_hybrid_indexer.go:concurrentBatchSave on
|
||||
// upstream/main):
|
||||
//
|
||||
// additionalParams["embedding"] map[string][]float32 // keyed by IndexInfo.SourceID
|
||||
// additionalParams["chunk_enabled"] map[string]bool // keyed by IndexInfo.ChunkID
|
||||
//
|
||||
// Both keys may be absent. Missing "embedding" → keyword-only path
|
||||
// (vector field omitted from the OpenSearch document, routed to the
|
||||
// dim-less keywords index). Missing "chunk_enabled" → IndexInfo's own
|
||||
// IsEnabled field is used as-is.
|
||||
|
||||
// Save indexes a single chunk. Idempotent: _id is set to chunk_id so
|
||||
// re-indexing the same chunk overwrites instead of creating a duplicate.
|
||||
func (r *Repository) Save(
|
||||
ctx context.Context,
|
||||
info *types.IndexInfo,
|
||||
params map[string]any,
|
||||
) error {
|
||||
emb := lookupEmbedding(params, info.SourceID)
|
||||
enabled := lookupChunkEnabled(params, info.ChunkID, info.IsEnabled)
|
||||
dim := len(emb)
|
||||
|
||||
var targetIndex string
|
||||
if dim > 0 {
|
||||
if err := r.ensureReady(ctx, dim); err != nil {
|
||||
return err
|
||||
}
|
||||
targetIndex = r.indexAlias(dim)
|
||||
} else {
|
||||
// Keyword-only path: no embedding supplied. Route to the
|
||||
// dim-less <base>_keywords index.
|
||||
if err := r.ensureKeywordsIndex(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
targetIndex = r.keywordsIndex()
|
||||
}
|
||||
|
||||
doc, err := json.Marshal(toDoc(info, emb, enabled))
|
||||
if err != nil {
|
||||
return fmt.Errorf("opensearch: marshal doc: %w", err)
|
||||
}
|
||||
req := osapi.IndexReq{
|
||||
Index: targetIndex,
|
||||
DocumentID: info.ChunkID, // _id = chunk_id (idempotent)
|
||||
Body: bytes.NewReader(doc),
|
||||
}
|
||||
resp, err := r.client.Index(ctx, req)
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchSave NDJSON-bulk-indexes chunks. Inspects per-item errors and
|
||||
// returns a typed error if ANY item fails.
|
||||
//
|
||||
// Dim-aware pre-marshal size cap: estimateBulkBodyBytes predicts NDJSON
|
||||
// size before allocating; over-cap callers get ErrBatchTooLarge instead
|
||||
// of a 50MB allocation that's then rejected. The service layer is
|
||||
// responsible for dim-aware sub-batching above this.
|
||||
func (r *Repository) BatchSave(
|
||||
ctx context.Context,
|
||||
infos []*types.IndexInfo,
|
||||
params map[string]any,
|
||||
) error {
|
||||
if len(infos) == 0 {
|
||||
return nil
|
||||
}
|
||||
emb, dim, err := extractBatchEmbeddings(params, infos)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Pre-marshal size + count check.
|
||||
if estimated := estimateBulkBodyBytes(len(infos), dim); estimated > 10*1024*1024 {
|
||||
return fmt.Errorf(
|
||||
"opensearch: estimated bulk body %dB exceeds 10MB cap (n=%d, dim=%d): %w",
|
||||
estimated, len(infos), dim, ErrBatchTooLarge,
|
||||
)
|
||||
}
|
||||
if len(infos) > 1000 {
|
||||
return fmt.Errorf("opensearch: bulk n=%d exceeds 1000-doc cap: %w",
|
||||
len(infos), ErrBatchTooLarge)
|
||||
}
|
||||
|
||||
// Route by dim: dim==0 means the entire batch is keyword-only
|
||||
// (concurrentBatchSaveNoEmbedding path), use the dim-less index.
|
||||
var alias string
|
||||
if dim == 0 {
|
||||
if err := r.ensureKeywordsIndex(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
alias = r.keywordsIndex()
|
||||
} else {
|
||||
if err := r.ensureReady(ctx, dim); err != nil {
|
||||
return err
|
||||
}
|
||||
alias = r.indexAlias(dim)
|
||||
}
|
||||
|
||||
body, err := buildBulkBody(alias, infos, emb, params)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req := osapi.BulkReq{Body: bytes.NewReader(body)}
|
||||
resp, err := r.client.Bulk(ctx, req)
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp != nil {
|
||||
defer drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
// Cap response body at 64MB (bulks of 1000 items are ~1-2MB success,
|
||||
// can balloon with full per-item errors).
|
||||
return inspectBulkResponse(io.LimitReader(resp.Inspect().Response.Body, 64<<20))
|
||||
}
|
||||
|
||||
// estimateBulkBodyBytes returns a rough NDJSON size in bytes. Slight
|
||||
// overestimate so the cap fires before real allocation. Per-item:
|
||||
// - action header `{"index":{"_index":"...","_id":"..."}}\n` ≈ 100B
|
||||
// - doc body: dim*5 (float ASCII) + ~700B metadata + content ≈ 1024B
|
||||
func estimateBulkBodyBytes(n, dim int) int {
|
||||
return n * (100 + dim*5 + 1024)
|
||||
}
|
||||
|
||||
// buildBulkBody constructs the NDJSON request body. emb[i] may be nil/
|
||||
// empty — toDoc handles the keyword-only case by omitting the
|
||||
// "embedding" field. chunk_enabled lookup also happens per-doc.
|
||||
func buildBulkBody(alias string, infos []*types.IndexInfo, emb [][]float32, params map[string]any) ([]byte, error) {
|
||||
var buf bytes.Buffer
|
||||
for i, info := range infos {
|
||||
action := map[string]any{
|
||||
"index": map[string]any{
|
||||
"_index": alias,
|
||||
"_id": info.ChunkID,
|
||||
},
|
||||
}
|
||||
actionJSON, err := json.Marshal(action)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opensearch: marshal bulk action %d: %w", i, err)
|
||||
}
|
||||
buf.Write(actionJSON)
|
||||
buf.WriteByte('\n')
|
||||
|
||||
enabled := lookupChunkEnabled(params, info.ChunkID, info.IsEnabled)
|
||||
docJSON, err := json.Marshal(toDoc(info, emb[i], enabled))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opensearch: marshal doc %d: %w", i, err)
|
||||
}
|
||||
buf.Write(docJSON)
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
// inspectBulkResponse parses OpenSearch bulk response, returns a typed
|
||||
// error if any item has a per-item .error. Surfaces up to 5 distinct
|
||||
// error messages + total count for operator diagnosis.
|
||||
//
|
||||
// Does NOT include det.Error.Reason in the wrapped public error
|
||||
// (Reason can contain document content fragments). Only ID + bounded
|
||||
// enum Type is included. Full per-item Reason goes to log.Debug.
|
||||
func inspectBulkResponse(body io.Reader) error {
|
||||
var r struct {
|
||||
Errors bool `json:"errors"`
|
||||
Items []map[string]struct {
|
||||
ID string `json:"_id"`
|
||||
Status int `json:"status"`
|
||||
Error *struct {
|
||||
Type string `json:"type"`
|
||||
Reason string `json:"reason"`
|
||||
} `json:"error,omitempty"`
|
||||
} `json:"items"`
|
||||
}
|
||||
if err := json.NewDecoder(body).Decode(&r); err != nil {
|
||||
return fmt.Errorf("opensearch: parse bulk response: %w", ErrTransport)
|
||||
}
|
||||
if !r.Errors {
|
||||
return nil
|
||||
}
|
||||
log := logger.GetLogger(context.Background())
|
||||
var msgs []string
|
||||
total := 0
|
||||
for _, item := range r.Items {
|
||||
for op, det := range item { // {"index": {...}}
|
||||
if det.Error == nil {
|
||||
continue
|
||||
}
|
||||
total++
|
||||
// Full reason → debug log only (may contain doc content).
|
||||
log.Debugf("[OpenSearch] bulk item err: op=%s id=%s type=%s reason=%s",
|
||||
op, det.ID, det.Error.Type, det.Error.Reason)
|
||||
if len(msgs) < 5 {
|
||||
// Public message: ID (operator-meaningful) + Type (bounded enum).
|
||||
msgs = append(msgs, fmt.Sprintf("[%s %s] %s", op, det.ID, det.Error.Type))
|
||||
}
|
||||
}
|
||||
}
|
||||
if total == 0 {
|
||||
return nil // shouldn't happen if r.Errors=true, but defensive
|
||||
}
|
||||
return fmt.Errorf("opensearch: bulk partial failure (%d items failed, first 5: %s): %w",
|
||||
total, strings.Join(msgs, "; "), ErrTransport)
|
||||
}
|
||||
|
||||
// toDoc projects IndexInfo + embedding to the OpenSearch document
|
||||
// shape. source_type is `int` (not stringified) — mapping declared as
|
||||
// `integer` field for stable cross-enum-extension behavior. Omits the
|
||||
// "embedding" field entirely when emb is empty (keyword-only doc, no
|
||||
// k-NN graph entry).
|
||||
func toDoc(info *types.IndexInfo, emb []float32, enabled bool) map[string]any {
|
||||
doc := map[string]any{
|
||||
"chunk_id": info.ChunkID,
|
||||
"knowledge_id": info.KnowledgeID,
|
||||
"knowledge_base_id": info.KnowledgeBaseID,
|
||||
"source_id": info.SourceID,
|
||||
"source_type": int(info.SourceType),
|
||||
"tag_id": info.TagID,
|
||||
"content": info.Content,
|
||||
"is_enabled": enabled,
|
||||
"is_recommended": info.IsRecommended,
|
||||
}
|
||||
if len(emb) > 0 {
|
||||
doc["embedding"] = emb
|
||||
}
|
||||
return doc
|
||||
}
|
||||
|
||||
// lookupEmbedding returns the embedding vector for info.SourceID.
|
||||
// Empty slice (not nil) indicates "no embedding provided" — caller
|
||||
// MUST check len(emb) before calling ensureReady.
|
||||
//
|
||||
// Defense in depth: tolerates the wrong shape rather than crash. The
|
||||
// service layer is responsible for getting params right; we degrade
|
||||
// to keyword-only for this doc and log a warning instead of panicking
|
||||
// downstream.
|
||||
func lookupEmbedding(params map[string]any, sourceID string) []float32 {
|
||||
if params == nil {
|
||||
return nil
|
||||
}
|
||||
raw, ok := params["embedding"]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
embMap, ok := raw.(map[string][]float32)
|
||||
if !ok {
|
||||
logger.GetLogger(context.Background()).Warnf(
|
||||
"[OpenSearch] additionalParams[\"embedding\"] is %T, want map[string][]float32", raw)
|
||||
return nil
|
||||
}
|
||||
return embMap[sourceID] // nil/empty if SourceID not in map
|
||||
}
|
||||
|
||||
// lookupChunkEnabled overlays the chunk_enabled map on top of
|
||||
// info.IsEnabled. Returns the overlay value if present, else the
|
||||
// default (info.IsEnabled).
|
||||
func lookupChunkEnabled(params map[string]any, chunkID string, def bool) bool {
|
||||
if params == nil {
|
||||
return def
|
||||
}
|
||||
raw, ok := params["chunk_enabled"]
|
||||
if !ok {
|
||||
return def
|
||||
}
|
||||
enMap, ok := raw.(map[string]bool)
|
||||
if !ok {
|
||||
return def
|
||||
}
|
||||
if v, present := enMap[chunkID]; present {
|
||||
return v
|
||||
}
|
||||
return def
|
||||
}
|
||||
|
||||
// extractBatchEmbeddings projects the params["embedding"] map onto the
|
||||
// info list, returning per-info embeddings (some may be empty) and the
|
||||
// dimension inferred from the first non-empty entry. Returns dim=0 if
|
||||
// no info has an embedding — caller routes to the keywords-only path.
|
||||
//
|
||||
// Defense in depth: all non-empty embeddings MUST share the same
|
||||
// dimension. Mixed-dim → ErrDimensionMismatch.
|
||||
func extractBatchEmbeddings(params map[string]any, infos []*types.IndexInfo) ([][]float32, int, error) {
|
||||
out := make([][]float32, len(infos))
|
||||
dim := 0
|
||||
for i, info := range infos {
|
||||
emb := lookupEmbedding(params, info.SourceID)
|
||||
out[i] = emb
|
||||
if len(emb) > 0 {
|
||||
if dim == 0 {
|
||||
dim = len(emb)
|
||||
} else if len(emb) != dim {
|
||||
return nil, 0, fmt.Errorf(
|
||||
"opensearch: embedding[%s] dim=%d != first non-empty dim=%d: %w",
|
||||
info.SourceID, len(emb), dim, ErrDimensionMismatch,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
return out, dim, nil
|
||||
}
|
||||
|
||||
// DeleteByChunkIDList runs _delete_by_query with a terms filter on
|
||||
// chunk_id. Sync only — large batches (>1000) are rejected with
|
||||
// ErrBatchTooLarge; the async task path that lifts the cap arrives in
|
||||
// a later change.
|
||||
//
|
||||
// Refresh policy: opensearch-go v4's typed Params.Refresh is *bool, so
|
||||
// the wire-level "wait_for" value documented in the OpenSearch REST
|
||||
// API is not directly expressible. Refresh:&true forces immediate
|
||||
// segment flush (read-your-writes guaranteed, but expensive at scale).
|
||||
// A follow-up can revisit via a custom raw request if telemetry shows
|
||||
// the cost is material at production batch frequencies.
|
||||
func (r *Repository) DeleteByChunkIDList(
|
||||
ctx context.Context, chunkIDs []string, dim int, _ string,
|
||||
) error {
|
||||
return r.deleteByList(ctx, chunkIDs, dim, "chunk_id")
|
||||
}
|
||||
|
||||
func (r *Repository) DeleteBySourceIDList(
|
||||
ctx context.Context, sourceIDs []string, dim int, _ string,
|
||||
) error {
|
||||
return r.deleteByList(ctx, sourceIDs, dim, "source_id")
|
||||
}
|
||||
|
||||
func (r *Repository) DeleteByKnowledgeIDList(
|
||||
ctx context.Context, knowledgeIDs []string, dim int, _ string,
|
||||
) error {
|
||||
return r.deleteByList(ctx, knowledgeIDs, dim, "knowledge_id")
|
||||
}
|
||||
|
||||
// deleteByList factors the common cap / empty / ensureReady / dispatch
|
||||
// logic out of the three DeleteBy* methods. dim==0 routes to the
|
||||
// dim-less keywords index.
|
||||
func (r *Repository) deleteByList(
|
||||
ctx context.Context, ids []string, dim int, field string,
|
||||
) error {
|
||||
if len(ids) == 0 {
|
||||
return nil
|
||||
}
|
||||
if len(ids) > 1000 {
|
||||
return fmt.Errorf("opensearch: %s-delete batch %d > 1000 cap: %w",
|
||||
field, len(ids), ErrBatchTooLarge)
|
||||
}
|
||||
var index string
|
||||
if dim == 0 {
|
||||
if err := r.ensureKeywordsIndex(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
index = r.keywordsIndex()
|
||||
} else {
|
||||
if err := r.ensureReady(ctx, dim); err != nil {
|
||||
return err
|
||||
}
|
||||
index = r.indexAlias(dim)
|
||||
}
|
||||
return r.deleteByTerms(ctx, index, field, ids)
|
||||
}
|
||||
|
||||
func (r *Repository) deleteByTerms(
|
||||
ctx context.Context, index, field string, values []string,
|
||||
) error {
|
||||
body, err := json.Marshal(map[string]any{
|
||||
"query": map[string]any{
|
||||
"terms": map[string]any{field: values},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("opensearch: marshal delete body: %w", err)
|
||||
}
|
||||
refresh := true
|
||||
req := osapi.DocumentDeleteByQueryReq{
|
||||
Indices: []string{index},
|
||||
Body: bytes.NewReader(body),
|
||||
Params: osapi.DocumentDeleteByQueryParams{Refresh: &refresh},
|
||||
}
|
||||
resp, err := r.client.Document.DeleteByQuery(ctx, req)
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
// Package opensearch implements the RetrieveEngineRepository interface
|
||||
// for OpenSearch k-NN native vector search.
|
||||
package opensearch
|
||||
|
||||
import "errors"
|
||||
|
||||
// Sentinel errors returned by Repository. The service-layer factory wraps
|
||||
// these into typed AppError values (2200/2201) — the repository itself
|
||||
// never imports internal/errors. The boundary is intentional (directional
|
||||
// dependency).
|
||||
var (
|
||||
// ErrIndexNotFound — alias / underlying index missing. Search and
|
||||
// delete-by-query operations return this when the per-dim alias has
|
||||
// not been created yet (no Save has been issued for that dim).
|
||||
ErrIndexNotFound = errors.New("opensearch: index not found")
|
||||
|
||||
// ErrDimensionMismatch — embedding dimension violates the per-dim
|
||||
// invariant (e.g. dim <= 0, dim > 16000, or embeddings within a
|
||||
// single batch disagree).
|
||||
ErrDimensionMismatch = errors.New("opensearch: embedding dimension mismatch")
|
||||
|
||||
// ErrAuth — cluster returned 401 / 403. Distinguished from ErrTransport
|
||||
// so the service layer can map to a clean 4xx instead of 503.
|
||||
ErrAuth = errors.New("opensearch: authentication failed")
|
||||
|
||||
// ErrTransport — network / 5xx / opaque cluster error. Classified as
|
||||
// transient: ensureReady does NOT persist this in initErr, so the next
|
||||
// caller will retry.
|
||||
ErrTransport = errors.New("opensearch: transport error")
|
||||
|
||||
// ErrVersionUnsupported — cluster is not OpenSearch, is OS 1.x, or is
|
||||
// OS 2.0~2.3 (pre-Lucene-HNSW-GA). probeVersion enforces.
|
||||
ErrVersionUnsupported = errors.New("opensearch: cluster version unsupported")
|
||||
|
||||
// ErrConfigInvalid — IndexConfig / storeID / sanitizeIndexName guard
|
||||
// failed, or the k-NN plugin is missing on one or more cluster nodes.
|
||||
ErrConfigInvalid = errors.New("opensearch: invalid index config")
|
||||
|
||||
// ErrFeatureNotEnabled — stubs.go returns this from methods whose real
|
||||
// implementation has not landed yet (CopyIndices / BatchUpdateChunk* /
|
||||
// swapToVersion, plus the read/write methods that a follow-up commit
|
||||
// will replace with production code).
|
||||
ErrFeatureNotEnabled = errors.New("opensearch: feature not enabled in this build")
|
||||
|
||||
// ErrBatchTooLarge — Save / Delete batch exceeded the driver's sync
|
||||
// cap. Distinct from ErrFeatureNotEnabled so the service layer can
|
||||
// chunk + retry rather than treat the failure as "waiting on a future
|
||||
// implementation."
|
||||
ErrBatchTooLarge = errors.New("opensearch: batch size exceeds driver cap")
|
||||
|
||||
// ErrCircuitBreaker — OpenSearch k-NN circuit breaker returned 429
|
||||
// (knn_circuit_breaker_exception). Classified as transient so the
|
||||
// caller can retry after the operator scales the cluster.
|
||||
ErrCircuitBreaker = errors.New("opensearch: knn circuit breaker open")
|
||||
)
|
||||
|
||||
// isTransientErr classifies sentinel errors. Transient errors do not get
|
||||
// persisted in ensureReady's initErr cache — they can be retried by the
|
||||
// next caller after the underlying cause clears.
|
||||
func isTransientErr(err error) bool {
|
||||
return errors.Is(err, ErrTransport) || errors.Is(err, ErrCircuitBreaker)
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// TestConnection verifies an OpenSearch cluster is reachable, runs a
|
||||
// supported version, and has the k-NN plugin installed on every node. It is
|
||||
// the connectivity probe used by the VectorStore service's CreateStore
|
||||
// health-check (the driver's unexported probes are reused here). Returns a
|
||||
// wrapped sentinel error on failure; nil on success.
|
||||
func TestConnection(ctx context.Context, cfg *types.ConnectionConfig) error {
|
||||
client, err := NewOpenSearchClient(cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := probeVersion(ctx, client); err != nil {
|
||||
return err
|
||||
}
|
||||
return probeKNNPlugin(ctx, client)
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// clusterHandler serves both GET / (version info) and /_cat/plugins, so the
|
||||
// full TestConnection probe (version + k-NN plugin) can run end to end.
|
||||
func clusterHandler(distribution, number string, plugins []osapi.CatPluginResp) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.URL.Path == "/_cat/plugins" {
|
||||
_ = json.NewEncoder(w).Encode(plugins)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"version":{"distribution":"` + distribution + `","number":"` + number + `"}}`))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTestConnection_Success(t *testing.T) {
|
||||
ts := httptest.NewServer(clusterHandler("opensearch", "3.3.2", []osapi.CatPluginResp{
|
||||
{Name: "node-1", Component: "opensearch-knn"},
|
||||
}))
|
||||
defer ts.Close()
|
||||
if err := TestConnection(context.Background(), &types.ConnectionConfig{Addr: ts.URL}); err != nil {
|
||||
t.Errorf("healthy cluster: want nil, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTestConnection_RejectsElasticsearch(t *testing.T) {
|
||||
ts := httptest.NewServer(clusterHandler("elasticsearch", "8.10.4", nil))
|
||||
defer ts.Close()
|
||||
if err := TestConnection(context.Background(), &types.ConnectionConfig{Addr: ts.URL}); err == nil {
|
||||
t.Error("elasticsearch cluster should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTestConnection_EmptyAddr(t *testing.T) {
|
||||
if err := TestConnection(context.Background(), &types.ConnectionConfig{}); err == nil {
|
||||
t.Error("empty addr should be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,420 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
)
|
||||
|
||||
// createIndexAndAlias creates "<alias>_v1" with the full k-NN mapping
|
||||
// and then aliases "<alias>" → "<alias>_v1". Idempotent: returns nil if
|
||||
// the alias already exists (so concurrent NewRepository calls in
|
||||
// different processes don't race the PUT).
|
||||
//
|
||||
// Orphan cleanup: if aliasPut fails after a successful indicesCreate,
|
||||
// the function attempts best-effort DELETE of the orphan _v1 index.
|
||||
// Cleanup failure is logged but does not mask the original aliasPut
|
||||
// error.
|
||||
//
|
||||
// Mapping-drift detection: if indicesCreate returns
|
||||
// resource_already_exists_exception (cross-process race), the existing
|
||||
// mapping's structural fingerprint is checked against the expected
|
||||
// fingerprint. Drift returns ErrConfigInvalid wrapping a clear "manual
|
||||
// reindex required" message instead of silently overlaying a foreign
|
||||
// mapping.
|
||||
func (r *Repository) createIndexAndAlias(ctx context.Context, dim int) error {
|
||||
log := logger.GetLogger(ctx)
|
||||
alias := r.indexAlias(dim)
|
||||
realIndex := alias + "_v1"
|
||||
|
||||
exists, err := r.aliasExists(ctx, alias)
|
||||
if err != nil {
|
||||
return fmt.Errorf("alias check %s: %w", alias, err)
|
||||
}
|
||||
if exists {
|
||||
return nil // another writer already set it up
|
||||
}
|
||||
|
||||
body, err := buildIndexMapping(r.cfg, dim)
|
||||
if err != nil {
|
||||
return fmt.Errorf("build mapping for dim=%d: %w", dim, err)
|
||||
}
|
||||
|
||||
indexCreated := false
|
||||
if err := r.indicesCreate(ctx, realIndex, body); err != nil {
|
||||
if isAlreadyExistsError(err) {
|
||||
// Cross-process race or stale orphan. Verify the existing
|
||||
// mapping matches what we expected; drift → ErrConfigInvalid.
|
||||
if driftErr := r.verifyMappingMatches(ctx, realIndex, body); driftErr != nil {
|
||||
return fmt.Errorf(
|
||||
"mapping drift on %s: existing index has incompatible mapping; "+
|
||||
"manual reindex required: %w",
|
||||
realIndex, driftErr,
|
||||
)
|
||||
}
|
||||
// Mapping matches → safe to proceed with alias setup.
|
||||
} else {
|
||||
return fmt.Errorf("create index %s: %w", realIndex, err)
|
||||
}
|
||||
} else {
|
||||
indexCreated = true
|
||||
}
|
||||
|
||||
if err := r.aliasPut(ctx, realIndex, alias); err != nil {
|
||||
// Best-effort orphan cleanup so future retries find a clean slate.
|
||||
if indexCreated {
|
||||
if delErr := r.indicesDelete(ctx, realIndex); delErr != nil {
|
||||
log.Warnf("[OpenSearch] orphan cleanup failed for %s: %v "+
|
||||
"(operator must DELETE manually)", realIndex, delErr)
|
||||
} else {
|
||||
log.Infof("[OpenSearch] cleaned up orphan %s after aliasPut failure", realIndex)
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("put alias %s → %s: %w", alias, realIndex, err)
|
||||
}
|
||||
if indexCreated {
|
||||
// Emit only when we actually provisioned the index (not when a
|
||||
// concurrent writer / existing index short-circuited above).
|
||||
r.auditSink().EmitIndexCreated(ctx, alias, dim)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// buildIndexMapping returns the full mapping JSON for a new OpenSearch
|
||||
// k-NN index. content uses the standard analyzer (BM25 default); all
|
||||
// *_id fields are explicit keyword sub-fields so DeleteByQuery / terms
|
||||
// filters hit them without the ES v8 "useKeywordSuffix" auto-detection
|
||||
// dance.
|
||||
//
|
||||
// match_type is intentionally NOT in the mapping — it's a per-result
|
||||
// classification derived from RetrieverType at parse time. source_type
|
||||
// is stored as `integer` (not stringified) so the wire format remains
|
||||
// stable across SourceType enum extension.
|
||||
//
|
||||
// Pure function — no IO, no logging. Easy to unit-test (byte-exact diff).
|
||||
func buildIndexMapping(cfg internalCfg, dim int) ([]byte, error) {
|
||||
type mapping struct {
|
||||
Settings struct {
|
||||
Index struct {
|
||||
Knn bool `json:"knn"`
|
||||
NumberOfShards int `json:"number_of_shards"`
|
||||
NumberOfReplicas int `json:"number_of_replicas"`
|
||||
RefreshInterval string `json:"refresh_interval"`
|
||||
KnnAlgoParamEFSearch int `json:"knn.algo_param.ef_search"`
|
||||
} `json:"index"`
|
||||
} `json:"settings"`
|
||||
Mappings struct {
|
||||
Properties map[string]any `json:"properties"`
|
||||
} `json:"mappings"`
|
||||
}
|
||||
var m mapping
|
||||
m.Settings.Index.Knn = true
|
||||
m.Settings.Index.NumberOfShards = cfg.shards
|
||||
m.Settings.Index.NumberOfReplicas = cfg.replicas
|
||||
m.Settings.Index.RefreshInterval = "1s"
|
||||
m.Settings.Index.KnnAlgoParamEFSearch = cfg.efSearch
|
||||
m.Mappings.Properties = map[string]any{
|
||||
"embedding": map[string]any{
|
||||
"type": "knn_vector",
|
||||
"dimension": dim,
|
||||
"method": map[string]any{
|
||||
"name": "hnsw",
|
||||
"space_type": "cosinesimil",
|
||||
"engine": cfg.knnEngine,
|
||||
"parameters": map[string]any{
|
||||
"m": cfg.hnswM,
|
||||
"ef_construction": cfg.hnswEFConstruction,
|
||||
},
|
||||
},
|
||||
},
|
||||
"content": map[string]any{"type": "text", "analyzer": "standard"},
|
||||
"chunk_id": map[string]any{"type": "keyword"},
|
||||
"knowledge_id": map[string]any{"type": "keyword"},
|
||||
"knowledge_base_id": map[string]any{"type": "keyword"},
|
||||
"tag_id": map[string]any{"type": "keyword"},
|
||||
"source_id": map[string]any{"type": "keyword"},
|
||||
"source_type": map[string]any{"type": "integer"},
|
||||
"is_enabled": map[string]any{"type": "boolean"},
|
||||
"is_recommended": map[string]any{"type": "boolean"},
|
||||
}
|
||||
out, err := json.Marshal(&m)
|
||||
if err != nil {
|
||||
// Cannot happen for our internal struct, but propagate for symmetry
|
||||
// with the other JSON builders in the package.
|
||||
return nil, fmt.Errorf("marshal mapping: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildKeywordsMapping returns the mapping for the dim-less keyword-only
|
||||
// index. Same as buildIndexMapping but with the embedding field omitted
|
||||
// — used by the no-embedding save path (the service layer dispatches
|
||||
// there when the knowledge base has vector indexing disabled).
|
||||
func buildKeywordsMapping(cfg internalCfg) ([]byte, error) {
|
||||
type mapping struct {
|
||||
Settings struct {
|
||||
Index struct {
|
||||
NumberOfShards int `json:"number_of_shards"`
|
||||
NumberOfReplicas int `json:"number_of_replicas"`
|
||||
RefreshInterval string `json:"refresh_interval"`
|
||||
} `json:"index"`
|
||||
} `json:"settings"`
|
||||
Mappings struct {
|
||||
Properties map[string]any `json:"properties"`
|
||||
} `json:"mappings"`
|
||||
}
|
||||
var m mapping
|
||||
m.Settings.Index.NumberOfShards = cfg.shards
|
||||
m.Settings.Index.NumberOfReplicas = cfg.replicas
|
||||
m.Settings.Index.RefreshInterval = "1s"
|
||||
m.Mappings.Properties = map[string]any{
|
||||
"content": map[string]any{"type": "text", "analyzer": "standard"},
|
||||
"chunk_id": map[string]any{"type": "keyword"},
|
||||
"knowledge_id": map[string]any{"type": "keyword"},
|
||||
"knowledge_base_id": map[string]any{"type": "keyword"},
|
||||
"tag_id": map[string]any{"type": "keyword"},
|
||||
"source_id": map[string]any{"type": "keyword"},
|
||||
"source_type": map[string]any{"type": "integer"},
|
||||
"is_enabled": map[string]any{"type": "boolean"},
|
||||
"is_recommended": map[string]any{"type": "boolean"},
|
||||
}
|
||||
out, err := json.Marshal(&m)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal keywords mapping: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ensureKeywordsIndex creates the dim-less keyword-only index used by
|
||||
// the no-embedding save path. mutex + flag pattern so transient
|
||||
// failures can be retried by the next caller (sync.Once cannot be
|
||||
// reset).
|
||||
func (r *Repository) ensureKeywordsIndex(ctx context.Context) error {
|
||||
r.keywordsMu.Lock()
|
||||
defer r.keywordsMu.Unlock()
|
||||
|
||||
if r.keywordsReady {
|
||||
return nil
|
||||
}
|
||||
if r.keywordsErr != nil && !isTransientErr(r.keywordsErr) {
|
||||
// Permanent failure — don't retry, surface the same error.
|
||||
return r.keywordsErr
|
||||
}
|
||||
|
||||
name := r.keywordsIndex()
|
||||
exists, err := r.aliasExists(ctx, name) // returns true if index exists too
|
||||
if err != nil {
|
||||
r.keywordsErr = err
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
r.keywordsReady = true
|
||||
r.keywordsErr = nil
|
||||
return nil
|
||||
}
|
||||
body, err := buildKeywordsMapping(r.cfg)
|
||||
if err != nil {
|
||||
r.keywordsErr = err
|
||||
return err
|
||||
}
|
||||
created := false
|
||||
if err := r.indicesCreate(ctx, name, body); err != nil {
|
||||
if !isAlreadyExistsError(err) {
|
||||
r.keywordsErr = err
|
||||
return err
|
||||
}
|
||||
// resource_already_exists_exception — race with concurrent process,
|
||||
// treat as success.
|
||||
} else {
|
||||
created = true
|
||||
}
|
||||
r.keywordsReady = true
|
||||
r.keywordsErr = nil
|
||||
if created {
|
||||
// dim=0 marks the dim-less keyword-only index.
|
||||
r.auditSink().EmitIndexCreated(ctx, name, 0)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// indicesCreate is the low-level wrapper around Indices.Create. Body is
|
||||
// the JSON mapping produced by buildIndexMapping.
|
||||
func (r *Repository) indicesCreate(ctx context.Context, name string, body []byte) error {
|
||||
req := osapi.IndicesCreateReq{Index: name, Body: bytes.NewReader(body)}
|
||||
resp, err := r.client.Indices.Create(ctx, req)
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// indicesDelete is used by createIndexAndAlias's orphan cleanup path.
|
||||
// Caller decides what to do with the error.
|
||||
func (r *Repository) indicesDelete(ctx context.Context, name string) error {
|
||||
req := osapi.IndicesDeleteReq{Indices: []string{name}}
|
||||
resp, err := r.client.Indices.Delete(ctx, req)
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// aliasExists uses HEAD /_alias/<name> (Indices.Alias.Exists) — returns
|
||||
// true on 200, false on 404, error on 5xx / transport.
|
||||
//
|
||||
// SDK quirk (opensearch-go v4.6.0): the Exists method passes
|
||||
// dataPointer=nil to its internal do(), which means non-2xx responses
|
||||
// come back as a plain *errors.errorString ("status: 404 Not Found")
|
||||
// rather than as a *opensearch.StructError. We therefore inspect
|
||||
// resp.StatusCode directly (resp is always returned even when err is
|
||||
// non-nil) and only fall back to wrapTransport for the "no response at
|
||||
// all" case (true network failure).
|
||||
func (r *Repository) aliasExists(ctx context.Context, alias string) (bool, error) {
|
||||
req := osapi.AliasExistsReq{Alias: []string{alias}}
|
||||
resp, err := r.client.Indices.Alias.Exists(ctx, req)
|
||||
if resp != nil {
|
||||
statusCode := resp.StatusCode
|
||||
drainAndClose(resp.Body)
|
||||
switch statusCode {
|
||||
case 200:
|
||||
return true, nil
|
||||
case 404:
|
||||
return false, nil
|
||||
}
|
||||
// Unexpected status (5xx) — fall through to the err path so
|
||||
// wrapTransport can classify (and so the test handler's
|
||||
// errBody is logged for diagnosis).
|
||||
}
|
||||
if err != nil {
|
||||
return false, wrapTransport(err)
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// aliasPut creates an alias pointing at index. Single add operation
|
||||
// (rolling-reindex atomic swap is deferred to a follow-up PR — see
|
||||
// stubs.go swapToVersion).
|
||||
func (r *Repository) aliasPut(ctx context.Context, index, alias string) error {
|
||||
req := osapi.AliasPutReq{
|
||||
Indices: []string{index},
|
||||
Alias: alias,
|
||||
}
|
||||
resp, err := r.client.Indices.Alias.Put(ctx, req)
|
||||
if err != nil {
|
||||
return wrapTransport(err)
|
||||
}
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// verifyMappingMatches GETs the existing index's mapping and compares
|
||||
// the embedding field's structural fingerprint (dimension, HNSW
|
||||
// parameters, engine) against the expected. A full field-by-field diff
|
||||
// would be ideal but expensive; we check only the parts that matter for
|
||||
// query correctness.
|
||||
func (r *Repository) verifyMappingMatches(ctx context.Context, index string, expected []byte) error {
|
||||
req := osapi.MappingGetReq{Indices: []string{index}}
|
||||
resp, err := r.client.Indices.Mapping.Get(ctx, &req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get mapping %s: %w", index, wrapTransport(err))
|
||||
}
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
|
||||
// Both fingerprints are derived from the mapping's embedding-field
|
||||
// HNSW parameters: dimension, m, ef_construction, engine, space_type.
|
||||
// We parse the expected (our own JSON) and the actual (cluster
|
||||
// response) and compare just those bits.
|
||||
expectedFP, err := extractEmbeddingFingerprint(expected)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse expected fingerprint: %w", err)
|
||||
}
|
||||
actualFP, err := extractMappingFingerprintFromResp(resp, index)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse actual fingerprint: %w", err)
|
||||
}
|
||||
if expectedFP != actualFP {
|
||||
return fmt.Errorf("mapping fingerprint mismatch: expected %+v, got %+v: %w",
|
||||
expectedFP, actualFP, ErrConfigInvalid)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// embFingerprint is the minimal set of HNSW parameters that uniquely
|
||||
// determine query semantics for a k-NN index. If any of these differ
|
||||
// between create-time and an existing index, retrieval is incorrect.
|
||||
type embFingerprint struct {
|
||||
Dimension int
|
||||
M int
|
||||
EFConstruction int
|
||||
Engine string
|
||||
SpaceType string
|
||||
}
|
||||
|
||||
// extractEmbeddingFingerprint parses our own mapping JSON (produced by
|
||||
// buildIndexMapping) to pull out the HNSW parameters.
|
||||
func extractEmbeddingFingerprint(body []byte) (embFingerprint, error) {
|
||||
var m struct {
|
||||
Mappings struct {
|
||||
Properties map[string]any `json:"properties"`
|
||||
} `json:"mappings"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &m); err != nil {
|
||||
return embFingerprint{}, err
|
||||
}
|
||||
return extractFingerprintFromProps(m.Mappings.Properties)
|
||||
}
|
||||
|
||||
// extractMappingFingerprintFromResp parses the cluster's GET-mapping
|
||||
// response and pulls out the embedding-field HNSW fingerprint for the
|
||||
// given index.
|
||||
func extractMappingFingerprintFromResp(resp *osapi.MappingGetResp, index string) (embFingerprint, error) {
|
||||
indexMapping, ok := resp.Indices[index]
|
||||
if !ok {
|
||||
return embFingerprint{}, fmt.Errorf("index %s not in mapping response", index)
|
||||
}
|
||||
// MappingGetResp.Indices[index].Mappings is a json.RawMessage in
|
||||
// the SDK — we decode the inner properties shape ourselves.
|
||||
var m struct {
|
||||
Properties map[string]any `json:"properties"`
|
||||
}
|
||||
if err := json.Unmarshal(indexMapping.Mappings, &m); err != nil {
|
||||
return embFingerprint{}, err
|
||||
}
|
||||
return extractFingerprintFromProps(m.Properties)
|
||||
}
|
||||
|
||||
func extractFingerprintFromProps(props map[string]any) (embFingerprint, error) {
|
||||
emb, ok := props["embedding"].(map[string]any)
|
||||
if !ok {
|
||||
return embFingerprint{}, fmt.Errorf("embedding field missing or wrong type")
|
||||
}
|
||||
dim, _ := emb["dimension"].(float64) // json.Unmarshal produces float64 for numbers
|
||||
method, _ := emb["method"].(map[string]any)
|
||||
engine, _ := method["engine"].(string)
|
||||
spaceType, _ := method["space_type"].(string)
|
||||
params, _ := method["parameters"].(map[string]any)
|
||||
m, _ := params["m"].(float64)
|
||||
efc, _ := params["ef_construction"].(float64)
|
||||
return embFingerprint{
|
||||
Dimension: int(dim),
|
||||
M: int(m),
|
||||
EFConstruction: int(efc),
|
||||
Engine: engine,
|
||||
SpaceType: spaceType,
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// retrieveFilters is the typed shape of OpenSearch query filter
|
||||
// clauses. It is the ONLY way callers (Retrieve) can constrain a
|
||||
// query — there is no JSON-fragment escape hatch. This prevents the
|
||||
// "map[string]interface{} injection" failure mode where a hostile
|
||||
// chunk_id like {"$or":[...]} would otherwise be inlined as JSON DSL.
|
||||
//
|
||||
// All slice fields are interpreted as IN-set membership (OS terms
|
||||
// query). Empty slice = no constraint on that dimension.
|
||||
//
|
||||
// IncludeDisabled controls the implicit is_enabled = true clause:
|
||||
// - false (default) — appends is_enabled term, hiding disabled chunks
|
||||
// - true — omits the clause; admin / audit callers can
|
||||
// observe disabled chunks. Service layer pins false for user-facing
|
||||
// APIs.
|
||||
type retrieveFilters struct {
|
||||
KBIDs []string
|
||||
KnowledgeIDs []string
|
||||
TagIDs []string
|
||||
ExcludeChunkIDs []string
|
||||
ExcludeKnowledgeIDs []string
|
||||
IncludeDisabled bool
|
||||
}
|
||||
|
||||
// fromParams projects RetrieveParams to retrieveFilters, applying the
|
||||
// multi-KB input from the service-layer dispatcher as terms.
|
||||
// ExcludeKnowledgeIDs is included so the OpenSearch driver matches the
|
||||
// behaviour of the Qdrant / pgvector drivers.
|
||||
func fromParams(p types.RetrieveParams) *retrieveFilters {
|
||||
return &retrieveFilters{
|
||||
KBIDs: p.KnowledgeBaseIDs,
|
||||
KnowledgeIDs: p.KnowledgeIDs,
|
||||
TagIDs: p.TagIDs,
|
||||
ExcludeChunkIDs: p.ExcludeChunkIDs,
|
||||
ExcludeKnowledgeIDs: p.ExcludeKnowledgeIDs,
|
||||
// IncludeDisabled stays false — set explicitly by admin callers
|
||||
// only. Driver receives this from a typed field, not from
|
||||
// AdditionalParams, so the contract is checked at compile time.
|
||||
}
|
||||
}
|
||||
|
||||
// toBoolMust emits the OpenSearch bool.must clause list for this
|
||||
// filter. Returns a freshly allocated slice each call so callers can
|
||||
// safely append additional clauses (e.g. the keyword match clause in
|
||||
// buildKeywordQuery) without aliasing.
|
||||
func (f *retrieveFilters) toBoolMust() []map[string]any {
|
||||
var must []map[string]any
|
||||
if len(f.KBIDs) > 0 {
|
||||
must = append(must, map[string]any{
|
||||
"terms": map[string]any{"knowledge_base_id": f.KBIDs},
|
||||
})
|
||||
}
|
||||
if len(f.KnowledgeIDs) > 0 {
|
||||
must = append(must, map[string]any{
|
||||
"terms": map[string]any{"knowledge_id": f.KnowledgeIDs},
|
||||
})
|
||||
}
|
||||
if len(f.TagIDs) > 0 {
|
||||
must = append(must, map[string]any{
|
||||
"terms": map[string]any{"tag_id": f.TagIDs},
|
||||
})
|
||||
}
|
||||
if len(f.ExcludeChunkIDs) > 0 {
|
||||
must = append(must, map[string]any{
|
||||
"bool": map[string]any{
|
||||
"must_not": map[string]any{
|
||||
"terms": map[string]any{"chunk_id": f.ExcludeChunkIDs},
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
if len(f.ExcludeKnowledgeIDs) > 0 {
|
||||
must = append(must, map[string]any{
|
||||
"bool": map[string]any{
|
||||
"must_not": map[string]any{
|
||||
"terms": map[string]any{"knowledge_id": f.ExcludeKnowledgeIDs},
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
if !f.IncludeDisabled {
|
||||
must = append(must, map[string]any{
|
||||
"term": map[string]any{"is_enabled": true},
|
||||
})
|
||||
}
|
||||
return must
|
||||
}
|
||||
|
||||
// buildKNNQuery composes a pure k-NN ANN query with caller filters. No
|
||||
// hybrid pipeline — native hybrid is out of scope for this PR; fusion
|
||||
// stays in the service-layer RRF.
|
||||
//
|
||||
// Applies `min_score` when threshold > 0. The OpenSearch k-NN plugin's
|
||||
// SpaceType.COSINESIMIL.scoreTranslation produces (1 + cos) / 2 ∈ [0, 1]
|
||||
// at the client, so the caller's Threshold value passes through to
|
||||
// min_score directly without further scaling.
|
||||
//
|
||||
// Signature returns ([]byte, error) so callers can propagate
|
||||
// json.Marshal failures (reachable for NaN / +/-Inf in embedding).
|
||||
func buildKNNQuery(embedding []float32, topK int, threshold float64, f *retrieveFilters) ([]byte, error) {
|
||||
knn := map[string]any{
|
||||
"embedding": map[string]any{
|
||||
"vector": embedding,
|
||||
"k": topK,
|
||||
"filter": map[string]any{
|
||||
"bool": map[string]any{"must": f.toBoolMust()},
|
||||
},
|
||||
},
|
||||
}
|
||||
body := map[string]any{
|
||||
"size": topK,
|
||||
"query": map[string]any{"knn": knn},
|
||||
}
|
||||
if threshold > 0 {
|
||||
body["min_score"] = threshold
|
||||
}
|
||||
out, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal knn query: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// buildKeywordQuery composes a BM25 match query against the content
|
||||
// field with caller filters. Score range is OpenSearch BM25 native (no
|
||||
// upper bound) — the normalizer keyword-passthrough rule treats keyword
|
||||
// retrievers as rank-based for RRF, so the score scale does not need
|
||||
// normalization here. min_score still applied if requested.
|
||||
func buildKeywordQuery(queryText string, topK int, threshold float64, f *retrieveFilters) ([]byte, error) {
|
||||
must := f.toBoolMust()
|
||||
must = append(must, map[string]any{
|
||||
"match": map[string]any{"content": queryText},
|
||||
})
|
||||
body := map[string]any{
|
||||
"size": topK,
|
||||
"query": map[string]any{"bool": map[string]any{"must": must}},
|
||||
}
|
||||
if threshold > 0 {
|
||||
body["min_score"] = threshold
|
||||
}
|
||||
out, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal keyword query: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@@ -0,0 +1,456 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/opensearch-project/opensearch-go/v4"
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
)
|
||||
|
||||
// EngineAwareNormalizer applies the documented per-engine cosine-score
|
||||
// formula. (Repository implements interfaces.RetrieveEngineRepository for
|
||||
// OpenSearch k-NN. Safe for concurrent use after construction: all fields
|
||||
// except the once / initErr maps are read-only after NewRepository returns;
|
||||
// per-dimension index initialization is guarded by sync.Once.
|
||||
//
|
||||
// Lifecycle: NewRepository validates connectivity + cluster version +
|
||||
// k-NN plugin, but does NOT create any index. Index creation happens
|
||||
// lazily on first Save / BatchSave / Retrieve once the embedding
|
||||
// dimension is known (see ensureReady — per-dimension index naming).
|
||||
//
|
||||
// Concurrency: A single Repository instance is shared across N goroutines
|
||||
// concurrently retrieving from the same store (the multi-store fan-out
|
||||
// dispatcher in the service layer). The struct is therefore designed to
|
||||
// be allocation-stable after construction; the only mutable state is the
|
||||
// sync.Once + initErr maps, each guarded by its own mutex.
|
||||
//
|
||||
// CRITICAL INVARIANT: initErr[dim] = err must be written INSIDE the
|
||||
// once.Do closure — moving it outside opens a TOCTOU where a concurrent
|
||||
// caller reads zero before the first caller persists.
|
||||
type Repository struct {
|
||||
client *osapi.Client
|
||||
baseIndex string // pre-suffix base name, e.g. "weknora_abcdef012345"
|
||||
cfg internalCfg // immutable after construction (value, not pointer)
|
||||
|
||||
// Per-dimension lazy initialization. ensureReady(ctx, dim) inserts an
|
||||
// entry the first time dim is seen, runs the create-index work inside
|
||||
// sync.Once, and persists any permanent error so subsequent callers
|
||||
// see the same failure without re-attempting the PUT. Transient
|
||||
// errors (ErrTransport, ErrCircuitBreaker) are NOT persisted — the
|
||||
// next caller retries.
|
||||
onceMu sync.Mutex
|
||||
once map[int]*sync.Once
|
||||
initMu sync.Mutex
|
||||
initErr map[int]error
|
||||
|
||||
// Lazy init state for the dim-less keyword-only index used by the
|
||||
// no-embedding path (concurrentBatchSaveNoEmbedding in the service
|
||||
// layer). Mutex + flag pattern (rather than sync.Once) so transient
|
||||
// failures can be retried by the next caller — sync.Once does not
|
||||
// support reset.
|
||||
keywordsMu sync.Mutex
|
||||
keywordsReady bool
|
||||
keywordsErr error
|
||||
|
||||
// sink receives audit events (index created / reindex executed). nil
|
||||
// means no auditing; use r.auditSink() to get a non-nil sink. Set via
|
||||
// WithAuditSink at construction.
|
||||
sink AuditSink
|
||||
}
|
||||
|
||||
// Compile-time interface satisfaction (Go best practice — keeps the build
|
||||
// red if the interface drifts and our implementation lags).
|
||||
var _ interfaces.RetrieveEngineRepository = (*Repository)(nil)
|
||||
|
||||
// NewRepository builds a new OpenSearch k-NN repository and verifies the
|
||||
// backing cluster is reachable + version-compatible + has the k-NN
|
||||
// plugin installed on every cluster node. It does NOT create any index —
|
||||
// callers (Save / Retrieve) trigger lazy per-dimension creation via
|
||||
// ensureReady on first use.
|
||||
//
|
||||
// storeID is the VectorStore.ID owning this repository instance. It is
|
||||
// folded into the base index name so multiple OpenSearch VectorStores
|
||||
// pointing at the same cluster do not collide:
|
||||
// - storeID == "" — env-path; no prefix, base index = "weknora".
|
||||
// - len(storeID) >= 16 — DB-store path; prefix = storeID[:12]
|
||||
// (48-bit collision space).
|
||||
// - 1..15 chars — rejected with ErrConfigInvalid.
|
||||
//
|
||||
// indexCfg is optional — pass nil to use env var (OPENSEARCH_INDEX) or
|
||||
// default ("weknora") values.
|
||||
//
|
||||
// Optional behavior is configured via functional options (e.g.
|
||||
// WithAuditSink). Passing no options keeps audit emission as a no-op, so the
|
||||
// env-path and tests need no extra wiring.
|
||||
//
|
||||
// Returns a typed sentinel error wrapped with %w; callers translate to
|
||||
// AppError at the engine-factory boundary.
|
||||
func NewRepository(
|
||||
ctx context.Context,
|
||||
client *osapi.Client,
|
||||
storeID string,
|
||||
indexCfg *types.IndexConfig,
|
||||
opts ...Option,
|
||||
) (interfaces.RetrieveEngineRepository, error) {
|
||||
log := logger.GetLogger(ctx)
|
||||
|
||||
if storeID != "" && len(storeID) < 16 {
|
||||
return nil, fmt.Errorf(
|
||||
"opensearch: storeID must be empty or >=16 chars, got %d: %w",
|
||||
len(storeID), ErrConfigInvalid,
|
||||
)
|
||||
}
|
||||
|
||||
base := types.ResolveIndexName(indexCfg, "OPENSEARCH_INDEX", "weknora")
|
||||
if storeID != "" {
|
||||
// 12 hex chars → 48-bit collision space (UUID is 32 chars; using
|
||||
// half gives strong-enough collision avoidance for the realistic
|
||||
// "two OpenSearch VectorStores on the same cluster" case).
|
||||
base = fmt.Sprintf("%s_%s", base, storeID[:12])
|
||||
}
|
||||
base, err := sanitizeIndexName(base)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opensearch: invalid index base name: %w", err)
|
||||
}
|
||||
|
||||
icfg, err := buildInternalCfg(indexCfg)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opensearch: invalid index config: %w", err)
|
||||
}
|
||||
|
||||
if err := probeVersion(ctx, client); err != nil {
|
||||
return nil, err // already wraps ErrVersionUnsupported / ErrTransport
|
||||
}
|
||||
if err := probeKNNPlugin(ctx, client); err != nil {
|
||||
return nil, err // already wraps ErrConfigInvalid / ErrTransport
|
||||
}
|
||||
|
||||
r := &Repository{
|
||||
client: client,
|
||||
baseIndex: base,
|
||||
cfg: icfg,
|
||||
once: make(map[int]*sync.Once),
|
||||
initErr: make(map[int]error),
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(r)
|
||||
}
|
||||
log.Infof("[OpenSearch] repository ready (baseIndex=%s, knn_engine=%s, hnsw_m=%d)",
|
||||
base, icfg.knnEngine, icfg.hnswM)
|
||||
return r, nil
|
||||
}
|
||||
|
||||
// ensureReady creates the per-dimension index (alias-backed) the first
|
||||
// time a given embedding dimension is seen. Concurrent callers for the
|
||||
// same dim block on the same sync.Once.
|
||||
//
|
||||
// Transient vs permanent error handling:
|
||||
// - PERMANENT failures (ErrConfigInvalid, ErrAuth, ErrVersionUnsupported,
|
||||
// ErrDimensionMismatch) are persisted in initErr[dim] and re-served to
|
||||
// all future callers without retry — caller must restart or fix config.
|
||||
// - TRANSIENT failures (ErrTransport, ErrCircuitBreaker) are NOT
|
||||
// persisted; the sync.Once is reset for that dim so the next caller
|
||||
// can retry. This prevents a single cluster blip from permanently
|
||||
// locking out a dim.
|
||||
//
|
||||
// Caller MUST pass a real ctx; nil is rejected. dim must be in (0, 16000]
|
||||
// (OpenSearch knn_vector hard limit).
|
||||
//
|
||||
// CRITICAL INVARIANT: initErr[dim] = err is written INSIDE the once.Do
|
||||
// closure. Moving it outside opens a TOCTOU where a concurrent caller
|
||||
// reads zero before the first caller persists.
|
||||
func (r *Repository) ensureReady(ctx context.Context, dim int) error {
|
||||
if ctx == nil {
|
||||
return fmt.Errorf("opensearch: ensureReady requires non-nil ctx")
|
||||
}
|
||||
if dim <= 0 || dim > 16000 {
|
||||
return fmt.Errorf("opensearch: dim %d out of range (1..16000): %w",
|
||||
dim, ErrDimensionMismatch)
|
||||
}
|
||||
|
||||
r.onceMu.Lock()
|
||||
once, ok := r.once[dim]
|
||||
if !ok {
|
||||
once = &sync.Once{}
|
||||
r.once[dim] = once
|
||||
}
|
||||
r.onceMu.Unlock()
|
||||
|
||||
once.Do(func() {
|
||||
err := r.createIndexAndAlias(ctx, dim)
|
||||
if err != nil && isTransientErr(err) {
|
||||
// Do NOT persist transient errors. Reset the sync.Once so
|
||||
// the next caller retries (requires once map mutation).
|
||||
r.onceMu.Lock()
|
||||
delete(r.once, dim)
|
||||
r.onceMu.Unlock()
|
||||
// Fall through — caller still sees this attempt's err.
|
||||
}
|
||||
r.initMu.Lock()
|
||||
if err == nil || !isTransientErr(err) {
|
||||
r.initErr[dim] = err // nil on success, persist permanent failures
|
||||
}
|
||||
r.initMu.Unlock()
|
||||
})
|
||||
|
||||
r.initMu.Lock()
|
||||
err := r.initErr[dim]
|
||||
r.initMu.Unlock()
|
||||
return err
|
||||
}
|
||||
|
||||
// indexAlias returns the alias name for a given dimension. All reads /
|
||||
// writes target the alias; the alias points at the actual <alias>_v1
|
||||
// index, enabling future rolling-reindex swap workflows (the swap path
|
||||
// itself is deferred to a follow-up PR; see stubs.go swapToVersion).
|
||||
func (r *Repository) indexAlias(dim int) string {
|
||||
return fmt.Sprintf("%s_%d", r.baseIndex, dim)
|
||||
}
|
||||
|
||||
// keywordsIndex returns the name of the dim-less keyword-only index used
|
||||
// by the no-embedding save path (the service layer dispatches there when
|
||||
// the knowledge base has vector indexing disabled). The mapping is the
|
||||
// same as the per-dim index minus the knn_vector field.
|
||||
func (r *Repository) keywordsIndex() string {
|
||||
return r.baseIndex + "_keywords"
|
||||
}
|
||||
|
||||
// EngineType returns the engine constant for OpenSearch. Used by the
|
||||
// fan-out normalizer to route to the [0, 1] passthrough group.
|
||||
func (r *Repository) EngineType() types.RetrieverEngineType {
|
||||
return types.OpenSearchRetrieverEngineType
|
||||
}
|
||||
|
||||
// Support reports the retrieval modes this driver implements. OpenSearch
|
||||
// k-NN supports both ANN (vector) and BM25 (keyword) within a single
|
||||
// document; the keyword-only index keeps BM25 functional for callers
|
||||
// that omit the embedding map.
|
||||
func (r *Repository) Support() []types.RetrieverType {
|
||||
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
|
||||
}
|
||||
|
||||
// probeVersion sends GET / and validates the cluster is OpenSearch in
|
||||
// a supported version range:
|
||||
// - Reject: ES (any), OS 1.x, OS 2.0~2.3 (Lucene HNSW preview).
|
||||
// - Warn-but-accept: OS 2.4~2.10 (Lucene HNSW GA, pre-LTS).
|
||||
// - Clean-accept: OS 2.11+, OS 3.x (primary tested: 3.3.2).
|
||||
func probeVersion(ctx context.Context, client *osapi.Client) error {
|
||||
resp, err := client.Info(ctx, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("opensearch: cluster info: %w", wrapTransport(err))
|
||||
}
|
||||
distribution := resp.Version.Distribution
|
||||
number := resp.Version.Number
|
||||
|
||||
if distribution != "opensearch" {
|
||||
return fmt.Errorf("opensearch: unsupported distribution %q: %w",
|
||||
distribution, ErrVersionUnsupported)
|
||||
}
|
||||
maj, min := parseMajorMinor(number)
|
||||
switch {
|
||||
case maj == 1:
|
||||
return fmt.Errorf("opensearch: 1.x EOL: %w", ErrVersionUnsupported)
|
||||
case maj == 2 && min >= 0 && min <= 3:
|
||||
return fmt.Errorf("opensearch: %s lacks Lucene HNSW GA (need 2.4+): %w",
|
||||
number, ErrVersionUnsupported)
|
||||
case maj == 2 && min >= 4 && min <= 10:
|
||||
logger.GetLogger(ctx).Warnf(
|
||||
"[OpenSearch] using pre-2.11 cluster %s; recommend 2.11+ LTS", number)
|
||||
return nil
|
||||
case maj == 2 || maj == 3:
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("opensearch: unsupported version %s: %w",
|
||||
number, ErrVersionUnsupported)
|
||||
}
|
||||
}
|
||||
|
||||
// parseMajorMinor extracts (major, minor) from a semver string such as
|
||||
// "3.3.2", "2.10.0-SNAPSHOT", "2.10". Returns (0, 0) for unparseable
|
||||
// input. Robust to pre-release suffixes (strips after '-') and missing
|
||||
// patch component (e.g. "2.10").
|
||||
func parseMajorMinor(num string) (major, minor int) {
|
||||
base := strings.SplitN(num, "-", 2)[0]
|
||||
parts := strings.Split(base, ".")
|
||||
if len(parts) < 2 {
|
||||
return 0, 0
|
||||
}
|
||||
m1, err1 := strconv.Atoi(parts[0])
|
||||
m2, err2 := strconv.Atoi(parts[1])
|
||||
if err1 != nil || err2 != nil {
|
||||
return 0, 0
|
||||
}
|
||||
return m1, m2
|
||||
}
|
||||
|
||||
// probeKNNPlugin verifies the OpenSearch cluster has the opensearch-knn
|
||||
// plugin installed on EVERY node. Without it on at least one node,
|
||||
// knn_vector queries that route a shard to that node fail with opaque
|
||||
// shard-failure errors — we surface this as a clear "plugin missing on
|
||||
// N nodes" message at registration.
|
||||
//
|
||||
// Implementation: _cat/plugins returns one row per (node, plugin) pair.
|
||||
// We group by node name and verify every distinct node has the
|
||||
// opensearch-knn component.
|
||||
func probeKNNPlugin(ctx context.Context, client *osapi.Client) error {
|
||||
// CatPluginsReq's Params struct has no Format field — the SDK
|
||||
// negotiates JSON via Accept header and returns a typed
|
||||
// CatPluginsResp with .Plugins []CatPluginResp already parsed.
|
||||
resp, err := client.Cat.Plugins(ctx, &osapi.CatPluginsReq{})
|
||||
if err != nil {
|
||||
return fmt.Errorf("opensearch: cat plugins: %w", wrapTransport(err))
|
||||
}
|
||||
|
||||
// Group by node, verify every distinct node has opensearch-knn.
|
||||
nodes := make(map[string]bool) // nodeName -> hasKNN
|
||||
for _, p := range resp.Plugins {
|
||||
if p.Name == "" {
|
||||
continue
|
||||
}
|
||||
if _, seen := nodes[p.Name]; !seen {
|
||||
nodes[p.Name] = false
|
||||
}
|
||||
if p.Component == "opensearch-knn" {
|
||||
nodes[p.Name] = true
|
||||
}
|
||||
}
|
||||
if len(nodes) == 0 {
|
||||
return fmt.Errorf("opensearch: _cat/plugins returned no rows: %w", ErrConfigInvalid)
|
||||
}
|
||||
var missing []string
|
||||
for node, hasKNN := range nodes {
|
||||
if !hasKNN {
|
||||
missing = append(missing, node)
|
||||
}
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
return fmt.Errorf(
|
||||
"opensearch: opensearch-knn plugin missing on %d/%d nodes (%v): %w",
|
||||
len(missing), len(nodes), missing, ErrConfigInvalid,
|
||||
)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// wrapTransport classifies a raw transport error into one of the
|
||||
// driver's sentinel errors based on HTTP status ONLY. The cluster-side
|
||||
// error message (err.Error()) is intentionally NOT included in the
|
||||
// wrapped public error — it may contain internal index names, shard
|
||||
// IDs, node hostnames, or document content snippets. The full SDK
|
||||
// error is routed to log.Debug for operator forensics.
|
||||
//
|
||||
// Classification:
|
||||
// - 401 / 403 → ErrAuth
|
||||
// - 404 → caller decides (use isNotFound separately for index ops)
|
||||
// - 429 + knn_circuit_breaker_exception → ErrCircuitBreaker
|
||||
// - 5xx / network / timeout → ErrTransport
|
||||
//
|
||||
// This is the ONE place in the package that maps wire errors to
|
||||
// sentinels; downstream functions just propagate.
|
||||
func wrapTransport(err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
// Route the verbose SDK message to debug log (operator-visible, not
|
||||
// user-visible). Background ctx is intentional — wrapTransport may
|
||||
// be invoked from any caller including the constructor.
|
||||
logger.GetLogger(context.Background()).Debugf("[OpenSearch] transport err: %v", err)
|
||||
|
||||
var se *opensearch.StructError
|
||||
if errors.As(err, &se) {
|
||||
switch se.Status {
|
||||
case http.StatusUnauthorized, http.StatusForbidden:
|
||||
return fmt.Errorf("authentication failed: %w", ErrAuth)
|
||||
case http.StatusTooManyRequests:
|
||||
if se.Err.Type == "knn_circuit_breaker_exception" {
|
||||
return fmt.Errorf("circuit breaker open: %w", ErrCircuitBreaker)
|
||||
}
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("transport error: %w", ErrTransport)
|
||||
}
|
||||
|
||||
// isNotFound returns true if err carries an HTTP 404. Used by index /
|
||||
// search code paths that want to distinguish "missing alias" from
|
||||
// "transport problem".
|
||||
func isNotFound(err error) bool {
|
||||
var se *opensearch.StructError
|
||||
return errors.As(err, &se) && se.Status == http.StatusNotFound
|
||||
}
|
||||
|
||||
// isAlreadyExistsError returns true if err is a 400 with OpenSearch's
|
||||
// resource_already_exists_exception type. Used by createIndexAndAlias
|
||||
// to handle the cross-process race where two replicas try to create
|
||||
// the same per-dim index simultaneously.
|
||||
func isAlreadyExistsError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var se *opensearch.StructError
|
||||
if !errors.As(err, &se) {
|
||||
return false
|
||||
}
|
||||
return se.Status == http.StatusBadRequest &&
|
||||
se.Err.Type == "resource_already_exists_exception"
|
||||
}
|
||||
|
||||
// sanitizeRe enforces the OpenSearch index-name contract:
|
||||
// - lowercase ASCII letters, digits, underscore, hyphen only
|
||||
// - must start with [a-z0-9]
|
||||
// - length 1..255
|
||||
var sanitizeRe = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,254}$`)
|
||||
|
||||
// sanitizeIndexName rejects names that violate the OpenSearch index-name
|
||||
// rules. Any non-conforming input returns ErrConfigInvalid — the driver
|
||||
// does NOT silently rewrite invalid input so operators get a clear
|
||||
// error at registration.
|
||||
func sanitizeIndexName(name string) (string, error) {
|
||||
if name == "" {
|
||||
return "", fmt.Errorf("empty index name: %w", ErrConfigInvalid)
|
||||
}
|
||||
// Explicit reject of OpenSearch-reserved wildcards / separators /
|
||||
// control characters. Some of these would already be caught by the
|
||||
// regexp below, but a dedicated check keeps the error message
|
||||
// specific.
|
||||
if strings.ContainsAny(name, "*?,\n\r\t/\\") {
|
||||
return "", fmt.Errorf("invalid char in %q: %w", name, ErrConfigInvalid)
|
||||
}
|
||||
if !sanitizeRe.MatchString(name) {
|
||||
return "", fmt.Errorf("name %q must match %s: %w",
|
||||
name, sanitizeRe.String(), ErrConfigInvalid)
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
|
||||
// drainAndClose closes the response body after draining unread bytes —
|
||||
// required for HTTP keep-alive connection reuse. opensearch-go does NOT
|
||||
// drain the body for us; failing to drain on every response will starve
|
||||
// the per-host connection pool under load.
|
||||
func drainAndClose(body io.ReadCloser) {
|
||||
if body == nil {
|
||||
return
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, body)
|
||||
_ = body.Close()
|
||||
}
|
||||
|
||||
// limitedDecode wraps json.NewDecoder with an io.LimitReader so a
|
||||
// pathological or malicious cluster response cannot exhaust process
|
||||
// memory. Callers pass the documented per-endpoint cap (e.g. 64MB for
|
||||
// bulk, 16MB for search, 1MB for plugin probe).
|
||||
func limitedDecode(body io.Reader, maxBytes int64, v any) error {
|
||||
return json.NewDecoder(io.LimitReader(body, maxBytes)).Decode(v)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,205 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// Retrieve dispatches to vectorRetrieve or keywordsRetrieve based on
|
||||
// params.RetrieverType. Returns exactly one *RetrieveResult — the
|
||||
// service-layer multi-store dispatcher depends on this one-bundle-per-
|
||||
// driver protocol. Score in result is RAW [0, 1] passthrough (the
|
||||
// OpenSearch k-NN plugin's SpaceType.COSINESIMIL.scoreTranslation has
|
||||
// already mapped to (1 + cos) / 2 before the score reaches us).
|
||||
//
|
||||
// Dim resolution order:
|
||||
// 1. params.AdditionalParams["dim"] if present and positive
|
||||
// 2. len(params.Embedding) if positive
|
||||
// 3. cross-dim multi-index search for keyword-only paths via
|
||||
// "<base>_*" wildcard pattern
|
||||
//
|
||||
// Vector retrieve always requires a positive dim (it's intrinsic to k-NN).
|
||||
func (r *Repository) Retrieve(
|
||||
ctx context.Context,
|
||||
params types.RetrieveParams,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
dim, multiIndex := resolveDim(params)
|
||||
|
||||
switch params.RetrieverType {
|
||||
case types.VectorRetrieverType:
|
||||
if dim == 0 {
|
||||
return nil, fmt.Errorf(
|
||||
"opensearch: vector retrieve requires embedding or AdditionalParams[\"dim\"]: %w",
|
||||
ErrDimensionMismatch)
|
||||
}
|
||||
if err := r.ensureReady(ctx, dim); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.vectorRetrieve(ctx, params, dim)
|
||||
case types.KeywordsRetrieverType:
|
||||
if multiIndex {
|
||||
// BM25-only path without dim — search across "<base>_*"
|
||||
// wildcard pattern. Slower than a single-index search but
|
||||
// correct for admin / keyword-only callers.
|
||||
return r.keywordsRetrieve(ctx, params, r.baseIndex+"_*")
|
||||
}
|
||||
if err := r.ensureReady(ctx, dim); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.keywordsRetrieve(ctx, params, r.indexAlias(dim))
|
||||
default:
|
||||
return nil, fmt.Errorf("opensearch: unsupported retriever type %q", params.RetrieverType)
|
||||
}
|
||||
}
|
||||
|
||||
// resolveDim returns (dim, multiIndex). multiIndex=true means dim is
|
||||
// unknown and the caller should use a wildcard alias pattern (keyword
|
||||
// path only — vector retrieve cannot be dim-less).
|
||||
func resolveDim(p types.RetrieveParams) (int, bool) {
|
||||
if p.AdditionalParams != nil {
|
||||
if v, ok := p.AdditionalParams["dim"].(int); ok && v > 0 {
|
||||
return v, false
|
||||
}
|
||||
}
|
||||
if len(p.Embedding) > 0 {
|
||||
return len(p.Embedding), false
|
||||
}
|
||||
return 0, true // unknown — multi-index fallback for keyword
|
||||
}
|
||||
|
||||
func (r *Repository) vectorRetrieve(
|
||||
ctx context.Context, p types.RetrieveParams, dim int,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
body, err := buildKNNQuery(p.Embedding, effectiveTopK(ctx, p), p.Threshold, fromParams(p))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hits, err := r.search(ctx, r.indexAlias(dim), body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return wrapResults(ctx, hits, types.VectorRetrieverType, types.MatchTypeEmbedding), nil
|
||||
}
|
||||
|
||||
func (r *Repository) keywordsRetrieve(
|
||||
ctx context.Context, p types.RetrieveParams, indexPattern string,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
body, err := buildKeywordQuery(p.Query, effectiveTopK(ctx, p), p.Threshold, fromParams(p))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hits, err := r.search(ctx, indexPattern, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return wrapResults(ctx, hits, types.KeywordsRetrieverType, types.MatchTypeKeywords), nil
|
||||
}
|
||||
|
||||
// effectiveTopK returns the caller's TopK or 10 if unset. The 10 is a
|
||||
// defensive default — every real caller sets TopK, so a TopK=10 in
|
||||
// production logs indicates a caller bug. Caps at 10000 to stay below
|
||||
// OpenSearch's index.max_result_window default.
|
||||
func effectiveTopK(ctx context.Context, p types.RetrieveParams) int {
|
||||
if p.TopK <= 0 {
|
||||
logger.GetLogger(ctx).Warnf("[OpenSearch] Retrieve called with TopK<=0; defaulting to 10 (caller bug?)")
|
||||
return 10
|
||||
}
|
||||
if p.TopK > 10000 {
|
||||
return 10000
|
||||
}
|
||||
return p.TopK
|
||||
}
|
||||
|
||||
func (r *Repository) search(ctx context.Context, indexPattern string, body []byte) ([]hit, error) {
|
||||
req := osapi.SearchReq{Indices: []string{indexPattern}, Body: bytes.NewReader(body)}
|
||||
resp, err := r.client.Search(ctx, &req)
|
||||
if err != nil {
|
||||
if isNotFound(err) {
|
||||
return nil, fmt.Errorf("opensearch: index %s missing: %w", indexPattern, ErrIndexNotFound)
|
||||
}
|
||||
return nil, wrapTransport(err)
|
||||
}
|
||||
defer func() {
|
||||
if resp != nil {
|
||||
drainAndClose(resp.Inspect().Response.Body)
|
||||
}
|
||||
}()
|
||||
// Cap response body at 16MB (search results bounded by max_result_window * doc_size).
|
||||
return parseSearchHits(io.LimitReader(resp.Inspect().Response.Body, 16<<20))
|
||||
}
|
||||
|
||||
// hit is the structural shape we pull from each OpenSearch search hit.
|
||||
// Field-by-field decode (vs map[string]any) keeps the JSON shape
|
||||
// pinned at compile time.
|
||||
type hit struct {
|
||||
ID string `json:"_id"` // equals chunk_id per the indexing invariant
|
||||
Score float64 `json:"_score"`
|
||||
Source struct {
|
||||
Content string `json:"content"`
|
||||
ChunkID string `json:"chunk_id"`
|
||||
KnowledgeID string `json:"knowledge_id"`
|
||||
KnowledgeBaseID string `json:"knowledge_base_id"`
|
||||
SourceID string `json:"source_id"`
|
||||
SourceType int `json:"source_type"` // integer, not stringified
|
||||
TagID string `json:"tag_id"`
|
||||
IsEnabled bool `json:"is_enabled"`
|
||||
} `json:"_source"`
|
||||
}
|
||||
|
||||
func parseSearchHits(body io.Reader) ([]hit, error) {
|
||||
var resp struct {
|
||||
Hits struct {
|
||||
Hits []hit `json:"hits"`
|
||||
} `json:"hits"`
|
||||
}
|
||||
if err := json.NewDecoder(body).Decode(&resp); err != nil {
|
||||
return nil, fmt.Errorf("opensearch: parse search response: %w", ErrTransport)
|
||||
}
|
||||
return resp.Hits.Hits, nil
|
||||
}
|
||||
|
||||
// wrapResults converts OpenSearch hits to the upstream IndexWithScore
|
||||
// shape. Always returns exactly one *RetrieveResult — the service-layer
|
||||
// multi-store dispatcher depends on this "one engine result per driver"
|
||||
// contract.
|
||||
//
|
||||
// Defensive: warns if hit._id != _source.chunk_id (we set _id=chunk_id
|
||||
// on every write, so a mismatch means external tooling has been writing
|
||||
// into the index) but does not fail the query — the result is still
|
||||
// usable for ranking, just suspicious.
|
||||
func wrapResults(ctx context.Context, hits []hit, rt types.RetrieverType, mt types.MatchType) []*types.RetrieveResult {
|
||||
log := logger.GetLogger(ctx)
|
||||
out := make([]*types.IndexWithScore, 0, len(hits))
|
||||
for _, h := range hits {
|
||||
if h.ID != h.Source.ChunkID {
|
||||
log.Warnf("[OpenSearch] hit._id=%q != _source.chunk_id=%q (D12 invariant violation)",
|
||||
h.ID, h.Source.ChunkID)
|
||||
}
|
||||
out = append(out, &types.IndexWithScore{
|
||||
ID: h.ID,
|
||||
ChunkID: h.Source.ChunkID,
|
||||
KnowledgeID: h.Source.KnowledgeID,
|
||||
KnowledgeBaseID: h.Source.KnowledgeBaseID,
|
||||
SourceID: h.Source.SourceID,
|
||||
SourceType: types.SourceType(h.Source.SourceType),
|
||||
TagID: h.Source.TagID,
|
||||
Content: h.Source.Content,
|
||||
Score: h.Score,
|
||||
MatchType: mt,
|
||||
IsEnabled: h.Source.IsEnabled,
|
||||
})
|
||||
}
|
||||
return []*types.RetrieveResult{{
|
||||
Results: out,
|
||||
RetrieverEngineType: types.OpenSearchRetrieverEngineType,
|
||||
RetrieverType: rt,
|
||||
}}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
// End-to-end repository tests intentionally use loopback httptest servers.
|
||||
secutils.SetSSRFWhitelistFromRaw("127.0.0.1,::1,localhost,opensearch.example.com")
|
||||
code := m.Run()
|
||||
secutils.SetSSRFWhitelistFromRaw("")
|
||||
os.Exit(code)
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// This file holds the remaining stubs for methods whose real implementation
|
||||
// has not landed yet — the rolling-reindex swap (swapToVersion) and the
|
||||
// precise storage estimate. CopyIndices (copy.go), BatchUpdateChunkEnabledStatus
|
||||
// and BatchUpdateChunkTagID (bulk_update.go) are now implemented. Each
|
||||
// remaining stub returns ErrFeatureNotEnabled (or, for EstimateStorageSize, a
|
||||
// conservative lower-bound) so any accidental invocation surfaces loudly.
|
||||
|
||||
// EstimateStorageSize: the real implementation that reads cluster
|
||||
// `_stats` for the per-dim alias arrives in a later change. For now we
|
||||
// return a conservative lower-bound estimate using the HNSW memory
|
||||
// formula so the upstream KB-delete guard fails-closed (treats non-empty
|
||||
// KBs as "may free non-trivial storage, force confirmation") rather than
|
||||
// failing open.
|
||||
//
|
||||
// Formula: N * (1024 content bytes + 4*dimGuess float + 128 HNSW M=16 overhead)
|
||||
// dimGuess = 768 (common embedding size; the real implementation reads
|
||||
// the actual dim from the cluster).
|
||||
func (r *Repository) EstimateStorageSize(
|
||||
_ context.Context,
|
||||
indexInfoList []*types.IndexInfo,
|
||||
_ map[string]any,
|
||||
) int64 {
|
||||
if len(indexInfoList) == 0 {
|
||||
return 0
|
||||
}
|
||||
const (
|
||||
contentBytes = 1024 // average chunk content
|
||||
embDimGuess = 768 // common embedding size
|
||||
hnswOverhead = 128 // M=16 → 8*M = 128 bytes/vector
|
||||
)
|
||||
return int64(len(indexInfoList)) * int64(contentBytes+4*embDimGuess+hnswOverhead)
|
||||
}
|
||||
|
||||
// swapToVersion is a stub for the future rolling-reindex swap path.
|
||||
// Calling it is illegal at this stage — alias is fixed at "_v1" and only
|
||||
// a follow-up change ships the swap orchestration. The surface exists so
|
||||
// the future API contract is reviewable now.
|
||||
//
|
||||
// Unexported because it is not part of the public
|
||||
// RetrieveEngineRepository interface — exposing it would require an
|
||||
// interface widening in internal/types/interfaces.
|
||||
func (r *Repository) swapToVersion(_ context.Context, _ int, _ int) error {
|
||||
return ErrFeatureNotEnabled
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/opensearch-project/opensearch-go/v4"
|
||||
osapi "github.com/opensearch-project/opensearch-go/v4/opensearchapi"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
secutils "github.com/Tencent/WeKnora/internal/utils"
|
||||
)
|
||||
|
||||
// NewOpenSearchClient builds a TLS-hardened, pool-tuned *osapi.Client for
|
||||
// the OpenSearch driver. The caller wires it into the registry from the
|
||||
// env path (container) and the DB-store path (engine factory); the
|
||||
// Repository constructor itself receives the pre-built client.
|
||||
//
|
||||
// TLS posture:
|
||||
// - MinVersion: TLS 1.2 (TLS 1.3 negotiated when both ends support).
|
||||
// - InsecureSkipVerify: opt-in via ConnectionConfig.InsecureSkipVerify.
|
||||
// Default false. Operators set true ONLY for self-signed dev clusters.
|
||||
// - CipherSuites: explicit forward-secrecy-only list. TLS 1.2's default
|
||||
// cipher list includes TLS_RSA_* (no forward secrecy); we drop them.
|
||||
// TLS 1.3 cipher selection is automatic and safe.
|
||||
//
|
||||
// Transport tuning:
|
||||
// - MaxIdleConnsPerHost: 32 (Go default 2 is too low when callers
|
||||
// dispatch concurrent queries across multiple knowledge bases)
|
||||
// - IdleConnTimeout: 90s (typical LB keep-alive)
|
||||
// - ResponseHeaderTimeout: 30s (per-request safety net)
|
||||
// - ExpectContinueTimeout: 1s
|
||||
func NewOpenSearchClient(cfg *types.ConnectionConfig) (*osapi.Client, error) {
|
||||
if cfg == nil || cfg.Addr == "" {
|
||||
return nil, fmt.Errorf("opensearch: ConnectionConfig.Addr required: %w", ErrConfigInvalid)
|
||||
}
|
||||
if err := secutils.ValidateURLForSSRF(cfg.Addr); err != nil {
|
||||
return nil, fmt.Errorf("opensearch: address failed SSRF validation: %w", err)
|
||||
}
|
||||
transport := buildHTTPTransport(cfg.InsecureSkipVerify)
|
||||
return osapi.NewClient(osapi.Config{
|
||||
Client: opensearch.Config{
|
||||
Addresses: []string{cfg.Addr},
|
||||
Username: cfg.Username,
|
||||
Password: cfg.Password,
|
||||
Transport: &secutils.SSRFValidatingRoundTripper{Base: transport},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// buildHTTPTransport constructs the TLS-hardened, pool-tuned
|
||||
// *http.Transport used by every OpenSearch client this driver creates.
|
||||
// Extracted from newOpenSearchClient so unit tests can pin the TLS
|
||||
// posture directly (the opensearch-go SDK wraps the transport in an
|
||||
// internal interface, making reflection-based introspection of the
|
||||
// final client fragile across SDK versions).
|
||||
func buildHTTPTransport(insecureSkipVerify bool) *http.Transport {
|
||||
return &http.Transport{
|
||||
Proxy: http.ProxyFromEnvironment,
|
||||
DialContext: secutils.SSRFSafeDialContext,
|
||||
TLSClientConfig: &tls.Config{
|
||||
MinVersion: tls.VersionTLS12,
|
||||
InsecureSkipVerify: insecureSkipVerify, //nolint:gosec — operator opt-in flag
|
||||
CipherSuites: []uint16{
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
|
||||
tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305,
|
||||
},
|
||||
},
|
||||
MaxIdleConnsPerHost: 32,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
ResponseHeaderTimeout: 30 * time.Second,
|
||||
ExpectContinueTimeout: 1 * time.Second,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package opensearch
|
||||
|
||||
import (
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
// TestNewOpenSearchClient_RejectsEmptyAddr verifies the constructor
|
||||
// refuses an empty Addr — operators who forget to set OPENSEARCH_ADDR
|
||||
// get a clear ErrConfigInvalid at registration instead of a cryptic
|
||||
// transport error later.
|
||||
func TestNewOpenSearchClient_RejectsEmptyAddr(t *testing.T) {
|
||||
t.Parallel()
|
||||
_, err := NewOpenSearchClient(&types.ConnectionConfig{Addr: ""})
|
||||
if !errors.Is(err, ErrConfigInvalid) {
|
||||
t.Fatalf("empty addr: want ErrConfigInvalid, got %v", err)
|
||||
}
|
||||
_, err = NewOpenSearchClient(nil)
|
||||
if !errors.Is(err, ErrConfigInvalid) {
|
||||
t.Fatalf("nil cfg: want ErrConfigInvalid, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewOpenSearchClient_Succeeds_OnValidAddr is a smoke test for the
|
||||
// happy path — a non-empty Addr must produce a non-nil *osapi.Client.
|
||||
// We don't probe the cluster (no network), just verify the constructor
|
||||
// returns successfully.
|
||||
func TestNewOpenSearchClient_Succeeds_OnValidAddr(t *testing.T) {
|
||||
t.Parallel()
|
||||
client, err := NewOpenSearchClient(&types.ConnectionConfig{
|
||||
Addr: "https://opensearch.example.com:9200",
|
||||
Username: "admin",
|
||||
Password: "secret", // not a real password — wire-format only
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewOpenSearchClient: %v", err)
|
||||
}
|
||||
if client == nil {
|
||||
t.Fatal("client must be non-nil on success")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildHTTPTransport_DefaultsToSecureTLS pins the default
|
||||
// InsecureSkipVerify=false. Operators who don't opt in MUST get a
|
||||
// strict-verify HTTPS client.
|
||||
func TestBuildHTTPTransport_DefaultsToSecureTLS(t *testing.T) {
|
||||
t.Parallel()
|
||||
tr := buildHTTPTransport(false)
|
||||
if tr.TLSClientConfig == nil {
|
||||
t.Fatal("TLSClientConfig must be set")
|
||||
}
|
||||
if tr.TLSClientConfig.InsecureSkipVerify {
|
||||
t.Fatal("default InsecureSkipVerify must be false (secure)")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildHTTPTransport_InsecureSkipVerifyOptIn confirms the flag
|
||||
// flows through to tls.Config. Self-signed dev clusters can opt in;
|
||||
// production deployments MUST leave it false.
|
||||
func TestBuildHTTPTransport_InsecureSkipVerifyOptIn(t *testing.T) {
|
||||
t.Parallel()
|
||||
tr := buildHTTPTransport(true)
|
||||
if tr.TLSClientConfig == nil {
|
||||
t.Fatal("TLSClientConfig must be set")
|
||||
}
|
||||
if !tr.TLSClientConfig.InsecureSkipVerify {
|
||||
t.Fatal("InsecureSkipVerify=true must propagate to tls.Config")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildHTTPTransport_TLS12MinVersion pins the TLS 1.2 hard floor.
|
||||
// TLS 1.3 negotiates automatically when both ends support it; 1.0 / 1.1
|
||||
// must be rejected at handshake time.
|
||||
func TestBuildHTTPTransport_TLS12MinVersion(t *testing.T) {
|
||||
t.Parallel()
|
||||
tr := buildHTTPTransport(false)
|
||||
if tr.TLSClientConfig.MinVersion != tls.VersionTLS12 {
|
||||
t.Fatalf("MinVersion: want tls.VersionTLS12 (0x%x), got 0x%x",
|
||||
tls.VersionTLS12, tr.TLSClientConfig.MinVersion)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildHTTPTransport_CipherSuitesExcludeNonECDHE pins the explicit
|
||||
// cipher list to forward-secrecy-only ECDHE suites. TLS 1.2's default
|
||||
// list includes TLS_RSA_* (no forward secrecy); we drop them. TLS 1.3
|
||||
// cipher selection is automatic so this only matters for the 1.2 path.
|
||||
func TestBuildHTTPTransport_CipherSuitesExcludeNonECDHE(t *testing.T) {
|
||||
t.Parallel()
|
||||
tr := buildHTTPTransport(false)
|
||||
cs := tr.TLSClientConfig.CipherSuites
|
||||
if len(cs) == 0 {
|
||||
t.Fatal("CipherSuites must be explicitly pinned (got empty)")
|
||||
}
|
||||
allowed := map[uint16]bool{
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: true,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256: true,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384: true,
|
||||
tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384: true,
|
||||
tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305: true,
|
||||
tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305: true,
|
||||
}
|
||||
for _, suite := range cs {
|
||||
if !allowed[suite] {
|
||||
t.Fatalf("cipher suite 0x%x is not in the forward-secrecy whitelist", suite)
|
||||
}
|
||||
}
|
||||
// Spot-check that a known-weak non-ECDHE cipher is NOT in the list.
|
||||
for _, suite := range cs {
|
||||
if suite == tls.TLS_RSA_WITH_AES_128_GCM_SHA256 {
|
||||
t.Fatal("TLS_RSA_WITH_AES_128_GCM_SHA256 must NOT be present (no forward secrecy)")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuildHTTPTransport_PoolTuning pins the connection-pool settings.
|
||||
// The Go default MaxIdleConnsPerHost is 2, which is too low for concurrent
|
||||
// queries across multiple knowledge bases.
|
||||
func TestBuildHTTPTransport_PoolTuning(t *testing.T) {
|
||||
t.Parallel()
|
||||
tr := buildHTTPTransport(false)
|
||||
if tr.MaxIdleConnsPerHost < 32 {
|
||||
t.Errorf("MaxIdleConnsPerHost: want >=32 for concurrent KB queries, got %d",
|
||||
tr.MaxIdleConnsPerHost)
|
||||
}
|
||||
if tr.IdleConnTimeout == 0 {
|
||||
t.Error("IdleConnTimeout must be set (Go default of 0 leaks idle conns)")
|
||||
}
|
||||
if tr.ResponseHeaderTimeout == 0 {
|
||||
t.Error("ResponseHeaderTimeout must be set (Go default of 0 means unbounded)")
|
||||
}
|
||||
if tr.ExpectContinueTimeout == 0 {
|
||||
t.Error("ExpectContinueTimeout must be set (avoids 100-continue stalls)")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,699 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/common"
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/google/uuid"
|
||||
"github.com/pgvector/pgvector-go"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// pgRepository implements PostgreSQL-based retrieval operations
|
||||
type pgRepository struct {
|
||||
db *gorm.DB // Database connection
|
||||
}
|
||||
|
||||
// NewPostgresRetrieveEngineRepository creates a new PostgreSQL retriever repository
|
||||
func NewPostgresRetrieveEngineRepository(db *gorm.DB) interfaces.RetrieveEngineRepository {
|
||||
logger.GetLogger(context.Background()).Info("[Postgres] Initializing PostgreSQL retriever engine repository")
|
||||
return &pgRepository{db: db}
|
||||
}
|
||||
|
||||
// EngineType returns the retriever engine type (PostgreSQL)
|
||||
func (r *pgRepository) EngineType() types.RetrieverEngineType {
|
||||
return types.PostgresRetrieverEngineType
|
||||
}
|
||||
|
||||
// Support returns supported retriever types (keywords and vector)
|
||||
func (r *pgRepository) Support() []types.RetrieverType {
|
||||
return []types.RetrieverType{types.KeywordsRetrieverType, types.VectorRetrieverType}
|
||||
}
|
||||
|
||||
// calculateIndexStorageSize calculates storage size for a single index entry
|
||||
func (g *pgRepository) calculateIndexStorageSize(embeddingDB *pgVector) int64 {
|
||||
// 1. Text content size
|
||||
contentSizeBytes := int64(len(embeddingDB.Content))
|
||||
|
||||
// 2. Vector storage size (2 bytes per dimension for half-precision float)
|
||||
var vectorSizeBytes int64 = 0
|
||||
if embeddingDB.Dimension > 0 {
|
||||
vectorSizeBytes = int64(embeddingDB.Dimension * 2)
|
||||
}
|
||||
|
||||
// 3. Metadata size (fixed overhead for IDs, timestamps etc.)
|
||||
metadataSizeBytes := int64(200)
|
||||
|
||||
// 4. Index overhead (HNSW index is ~2x vector size)
|
||||
indexOverheadBytes := vectorSizeBytes * 2
|
||||
|
||||
// Total size in bytes
|
||||
totalSizeBytes := contentSizeBytes + vectorSizeBytes + metadataSizeBytes + indexOverheadBytes
|
||||
|
||||
return totalSizeBytes
|
||||
}
|
||||
|
||||
// EstimateStorageSize estimates total storage size for multiple indices
|
||||
func (g *pgRepository) EstimateStorageSize(
|
||||
ctx context.Context, indexInfoList []*types.IndexInfo, additionalParams map[string]any,
|
||||
) int64 {
|
||||
var totalStorageSize int64 = 0
|
||||
for _, indexInfo := range indexInfoList {
|
||||
embeddingDB := toDBVectorEmbedding(indexInfo, additionalParams)
|
||||
totalStorageSize += g.calculateIndexStorageSize(embeddingDB)
|
||||
}
|
||||
logger.GetLogger(ctx).Infof(
|
||||
"[Postgres] Estimated storage size for %d indices: %d bytes",
|
||||
len(indexInfoList), totalStorageSize,
|
||||
)
|
||||
return totalStorageSize
|
||||
}
|
||||
|
||||
// Save stores a single index entry
|
||||
func (g *pgRepository) Save(ctx context.Context, indexInfo *types.IndexInfo, additionalParams map[string]any) error {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Saving index for source ID: %s", indexInfo.SourceID)
|
||||
embeddingDB := toDBVectorEmbedding(indexInfo, additionalParams)
|
||||
err := g.db.WithContext(ctx).Create(embeddingDB).Error
|
||||
if err != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to save index: %v", err)
|
||||
return err
|
||||
}
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully saved index for source ID: %s", indexInfo.SourceID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchSave stores multiple index entries in batch
|
||||
func (g *pgRepository) BatchSave(
|
||||
ctx context.Context, indexInfoList []*types.IndexInfo, additionalParams map[string]any,
|
||||
) error {
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Batch saving %d indices", len(indexInfoList))
|
||||
indexInfoDBList := make([]*pgVector, len(indexInfoList))
|
||||
for i := range indexInfoList {
|
||||
indexInfoDBList[i] = toDBVectorEmbedding(indexInfoList[i], additionalParams)
|
||||
}
|
||||
err := g.db.WithContext(ctx).Clauses(clause.OnConflict{DoNothing: true}).Create(indexInfoDBList).Error
|
||||
if err != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Batch save failed: %v", err)
|
||||
return err
|
||||
}
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully batch saved %d indices", len(indexInfoList))
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteByChunkIDList deletes indices by chunk IDs
|
||||
func (g *pgRepository) DeleteByChunkIDList(ctx context.Context, chunkIDList []string, dimension int, knowledgeType string) error {
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Deleting indices by chunk IDs, count: %d", len(chunkIDList))
|
||||
result := g.db.WithContext(ctx).Where("chunk_id IN ?", chunkIDList).Delete(&pgVector{})
|
||||
if result.Error != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to delete indices by chunk IDs: %v", result.Error)
|
||||
return result.Error
|
||||
}
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully deleted %d indices by chunk IDs", result.RowsAffected)
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteBySourceIDList deletes indices by source IDs
|
||||
func (g *pgRepository) DeleteBySourceIDList(ctx context.Context, sourceIDList []string, dimension int, knowledgeType string) error {
|
||||
if len(sourceIDList) == 0 {
|
||||
return nil
|
||||
}
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Deleting indices by source IDs, count: %d", len(sourceIDList))
|
||||
result := g.db.WithContext(ctx).Where("source_id IN ?", sourceIDList).Delete(&pgVector{})
|
||||
if result.Error != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to delete indices by source IDs: %v", result.Error)
|
||||
return result.Error
|
||||
}
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully deleted %d indices by source IDs", result.RowsAffected)
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteByKnowledgeIDList deletes indices by knowledge IDs
|
||||
func (g *pgRepository) DeleteByKnowledgeIDList(ctx context.Context, knowledgeIDList []string, dimension int, knowledgeType string) error {
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Deleting indices by knowledge IDs, count: %d", len(knowledgeIDList))
|
||||
result := g.db.WithContext(ctx).Where("knowledge_id IN ?", knowledgeIDList).Delete(&pgVector{})
|
||||
if result.Error != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to delete indices by knowledge IDs: %v", result.Error)
|
||||
return result.Error
|
||||
}
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully deleted %d indices by knowledge IDs", result.RowsAffected)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Retrieve handles retrieval requests and routes to appropriate method
|
||||
func (g *pgRepository) Retrieve(ctx context.Context, params types.RetrieveParams) ([]*types.RetrieveResult, error) {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Processing retrieval request of type: %s", params.RetrieverType)
|
||||
switch params.RetrieverType {
|
||||
case types.KeywordsRetrieverType:
|
||||
return g.KeywordsRetrieve(ctx, params)
|
||||
case types.VectorRetrieverType:
|
||||
return g.VectorRetrieve(ctx, params)
|
||||
}
|
||||
err := errors.New("invalid retriever type")
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] %v: %s", err, params.RetrieverType)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// KeywordsRetrieve performs keyword-based search using PostgreSQL full-text search
|
||||
func (g *pgRepository) KeywordsRetrieve(ctx context.Context,
|
||||
params types.RetrieveParams,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Keywords retrieval: query=%s, topK=%d", params.Query, params.TopK)
|
||||
conds := make([]clause.Expression, 0)
|
||||
|
||||
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
|
||||
// - If only KnowledgeBaseIDs: search entire knowledge bases
|
||||
// - If only KnowledgeIDs: search specific documents
|
||||
// - If both: search specific documents within the knowledge bases (AND)
|
||||
if len(params.KnowledgeBaseIDs) > 0 {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Filtering by knowledge base IDs: %v", params.KnowledgeBaseIDs)
|
||||
conds = append(conds, clause.IN{
|
||||
Column: "knowledge_base_id",
|
||||
Values: common.ToInterfaceSlice(params.KnowledgeBaseIDs),
|
||||
})
|
||||
}
|
||||
if len(params.KnowledgeIDs) > 0 {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Filtering by knowledge IDs: %v", params.KnowledgeIDs)
|
||||
conds = append(conds, clause.IN{
|
||||
Column: "knowledge_id",
|
||||
Values: common.ToInterfaceSlice(params.KnowledgeIDs),
|
||||
})
|
||||
}
|
||||
// Filter by tag IDs if specified
|
||||
if len(params.TagIDs) > 0 {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Filtering by tag IDs: %v", params.TagIDs)
|
||||
conds = append(conds, clause.IN{
|
||||
Column: "tag_id",
|
||||
Values: common.ToInterfaceSlice(params.TagIDs),
|
||||
})
|
||||
}
|
||||
|
||||
// Use ParadeDB's ||| operator for matching any token
|
||||
conds = append(conds, clause.Expr{
|
||||
SQL: "content ||| ?",
|
||||
Vars: []interface{}{params.Query},
|
||||
})
|
||||
|
||||
// Filter by is_enabled = true or NULL (NULL means enabled for historical data)
|
||||
conds = append(conds, clause.Expr{
|
||||
SQL: "(is_enabled IS NULL OR is_enabled = ?)",
|
||||
Vars: []interface{}{true},
|
||||
})
|
||||
conds = append(conds, clause.OrderBy{Columns: []clause.OrderByColumn{
|
||||
{Column: clause.Column{Name: "score"}, Desc: true},
|
||||
}})
|
||||
|
||||
var embeddingDBList []pgVectorWithScore
|
||||
err := g.db.WithContext(ctx).Clauses(conds...).Debug().
|
||||
Select([]string{
|
||||
"paradedb.score(id) as score",
|
||||
"id",
|
||||
"content",
|
||||
"source_id",
|
||||
"source_type",
|
||||
"chunk_id",
|
||||
"knowledge_id",
|
||||
"knowledge_base_id",
|
||||
"tag_id",
|
||||
}).
|
||||
Limit(int(params.TopK)).
|
||||
Find(&embeddingDBList).Error
|
||||
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] No records found for keywords query: %s", params.Query)
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Keywords retrieval failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Keywords retrieval found %d results", len(embeddingDBList))
|
||||
results := make([]*types.IndexWithScore, len(embeddingDBList))
|
||||
const maxKeywordResultLog = 8
|
||||
for i := range embeddingDBList {
|
||||
results[i] = fromDBVectorEmbeddingWithScore(&embeddingDBList[i], types.MatchTypeKeywords)
|
||||
if i < maxKeywordResultLog {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Keywords result %d: chunk=%s, score=%f",
|
||||
i, results[i].ChunkID, results[i].Score)
|
||||
}
|
||||
}
|
||||
if len(results) > maxKeywordResultLog {
|
||||
logger.GetLogger(ctx).Debugf(
|
||||
"[Postgres] Keywords result summary: total=%d logged=%d truncated=%d",
|
||||
len(results), maxKeywordResultLog, len(results)-maxKeywordResultLog,
|
||||
)
|
||||
}
|
||||
return []*types.RetrieveResult{
|
||||
{
|
||||
Results: results,
|
||||
RetrieverEngineType: types.PostgresRetrieverEngineType,
|
||||
RetrieverType: types.KeywordsRetrieverType,
|
||||
Error: nil,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// VectorRetrieve performs vector similarity search using pgvector
|
||||
// Optimized to use HNSW index efficiently and avoid recalculating vector distance
|
||||
func (g *pgRepository) VectorRetrieve(ctx context.Context,
|
||||
params types.RetrieveParams,
|
||||
) ([]*types.RetrieveResult, error) {
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Vector retrieval: dim=%d, topK=%d, threshold=%.4f",
|
||||
len(params.Embedding), params.TopK, params.Threshold)
|
||||
|
||||
dimension := len(params.Embedding)
|
||||
queryVector := pgvector.NewHalfVector(params.Embedding)
|
||||
|
||||
// Build WHERE conditions for filtering
|
||||
whereParts := make([]string, 0)
|
||||
allVars := make([]interface{}, 0)
|
||||
|
||||
// Add query vector first (used in ORDER BY for HNSW index)
|
||||
allVars = append(allVars, queryVector)
|
||||
|
||||
// Dimension filter (required for HNSW index WHERE clause)
|
||||
whereParts = append(whereParts, fmt.Sprintf("dimension = $%d", len(allVars)+1))
|
||||
allVars = append(allVars, dimension)
|
||||
|
||||
// KnowledgeBaseIDs and KnowledgeIDs use AND logic
|
||||
// - If only KnowledgeBaseIDs: search entire knowledge bases
|
||||
// - If only KnowledgeIDs: search specific documents
|
||||
// - If both: search specific documents within the knowledge bases (AND)
|
||||
if len(params.KnowledgeBaseIDs) > 0 {
|
||||
logger.GetLogger(ctx).Debugf(
|
||||
"[Postgres] Filtering vector search by knowledge base IDs: %v",
|
||||
params.KnowledgeBaseIDs,
|
||||
)
|
||||
placeholders := make([]string, len(params.KnowledgeBaseIDs))
|
||||
paramStart := len(allVars) + 1
|
||||
for i := range params.KnowledgeBaseIDs {
|
||||
placeholders[i] = fmt.Sprintf("$%d", paramStart+i)
|
||||
allVars = append(allVars, params.KnowledgeBaseIDs[i])
|
||||
}
|
||||
whereParts = append(whereParts, fmt.Sprintf("knowledge_base_id IN (%s)",
|
||||
strings.Join(placeholders, ", ")))
|
||||
}
|
||||
if len(params.KnowledgeIDs) > 0 {
|
||||
logger.GetLogger(ctx).Debugf(
|
||||
"[Postgres] Filtering vector search by knowledge IDs: %v",
|
||||
params.KnowledgeIDs,
|
||||
)
|
||||
placeholders := make([]string, len(params.KnowledgeIDs))
|
||||
paramStart := len(allVars) + 1
|
||||
for i := range params.KnowledgeIDs {
|
||||
placeholders[i] = fmt.Sprintf("$%d", paramStart+i)
|
||||
allVars = append(allVars, params.KnowledgeIDs[i])
|
||||
}
|
||||
whereParts = append(whereParts, fmt.Sprintf("knowledge_id IN (%s)",
|
||||
strings.Join(placeholders, ", ")))
|
||||
}
|
||||
// Filter by tag IDs if specified
|
||||
if len(params.TagIDs) > 0 {
|
||||
logger.GetLogger(ctx).Debugf(
|
||||
"[Postgres] Filtering vector search by tag IDs: %v",
|
||||
params.TagIDs,
|
||||
)
|
||||
placeholders := make([]string, len(params.TagIDs))
|
||||
paramStart := len(allVars) + 1
|
||||
for i := range params.TagIDs {
|
||||
placeholders[i] = fmt.Sprintf("$%d", paramStart+i)
|
||||
allVars = append(allVars, params.TagIDs[i])
|
||||
}
|
||||
whereParts = append(whereParts, fmt.Sprintf("tag_id IN (%s)",
|
||||
strings.Join(placeholders, ", ")))
|
||||
}
|
||||
|
||||
// is_enabled filter
|
||||
whereParts = append(whereParts, fmt.Sprintf("(is_enabled IS NULL OR is_enabled = $%d)", len(allVars)+1))
|
||||
allVars = append(allVars, true)
|
||||
|
||||
// Build WHERE clause string
|
||||
whereClause := ""
|
||||
if len(whereParts) > 0 {
|
||||
whereClause = "WHERE " + strings.Join(whereParts, " AND ")
|
||||
}
|
||||
|
||||
// Expand TopK to get more candidates before threshold filtering.
|
||||
//
|
||||
// HNSW requires `ef_search >= LIMIT`, and a very large LIMIT (e.g. 1000)
|
||||
// forces HNSW to walk a near-exhaustive portion of the graph, often making
|
||||
// it slower than a sequential scan and pushing the planner to pick Seq Scan
|
||||
// even when an index exists. 200 is a good sweet spot: it gives enough
|
||||
// headroom for threshold/filter post-processing without ballooning ef_search.
|
||||
expandedTopK := params.TopK * 2
|
||||
if expandedTopK < 100 {
|
||||
expandedTopK = 100 // Minimum 100 candidates
|
||||
}
|
||||
if expandedTopK > 200 {
|
||||
expandedTopK = 200 // Maximum 200 candidates (keeps HNSW efficient)
|
||||
}
|
||||
if expandedTopK < params.TopK {
|
||||
expandedTopK = params.TopK // Ensure subquery limit is at least final limit
|
||||
}
|
||||
|
||||
// Optimized query: Use subquery to calculate distance once.
|
||||
//
|
||||
// IMPORTANT: The HNSW index in this project is built on the EXPRESSION
|
||||
// (embedding::halfvec(<dim>)) halfvec_cosine_ops
|
||||
// because the `embedding` column itself is `halfvec` without a fixed dimension
|
||||
// (so the table can store multiple embedding sizes such as 798 / 3584 / ...).
|
||||
//
|
||||
// pgvector requires the ORDER BY expression to match the indexed expression
|
||||
// EXACTLY, otherwise the planner falls back to a sequential scan. The
|
||||
// `embedding::halfvec(%d)` cast on both sides of `<=>` is therefore NOT
|
||||
// redundant — it is the only way to make the HNSW index get used at all.
|
||||
// See: pgvector issues #702, #835 and ParadeDB "indexing-expressions" docs.
|
||||
subqueryLimitParam := len(allVars) + 1
|
||||
thresholdParam := len(allVars) + 2
|
||||
finalLimitParam := len(allVars) + 3
|
||||
|
||||
querySQL := fmt.Sprintf(`
|
||||
SELECT
|
||||
id, content, source_id, source_type, chunk_id, knowledge_id, knowledge_base_id, tag_id,
|
||||
(1 - distance) as score
|
||||
FROM (
|
||||
SELECT
|
||||
id, content, source_id, source_type, chunk_id, knowledge_id, knowledge_base_id, tag_id,
|
||||
embedding::halfvec(%[1]d) <=> $1::halfvec(%[1]d) as distance
|
||||
FROM embeddings
|
||||
%[2]s
|
||||
ORDER BY embedding::halfvec(%[1]d) <=> $1::halfvec(%[1]d)
|
||||
LIMIT $%[3]d
|
||||
) AS candidates
|
||||
WHERE distance <= $%[4]d
|
||||
ORDER BY distance ASC
|
||||
LIMIT $%[5]d
|
||||
`, dimension, whereClause, subqueryLimitParam, thresholdParam, finalLimitParam)
|
||||
|
||||
allVars = append(allVars, expandedTopK) // LIMIT in subquery
|
||||
allVars = append(allVars, 1-params.Threshold) // Distance threshold
|
||||
allVars = append(allVars, params.TopK) // Final LIMIT
|
||||
|
||||
// HNSW's `ef_search` defaults to 40, which is much smaller than our
|
||||
// `expandedTopK` budget (up to 1000). Without raising it, HNSW would only
|
||||
// return ~40 candidates per scan and our outer threshold/filter step would
|
||||
// silently lose recall. `SET LOCAL` requires a transaction so we wrap the
|
||||
// query in one. We never write inside this transaction, so the cost is
|
||||
// negligible.
|
||||
efSearch := expandedTopK
|
||||
if efSearch < 40 {
|
||||
efSearch = 40
|
||||
}
|
||||
|
||||
var embeddingDBList []pgVectorWithScore
|
||||
|
||||
err := g.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
if err := tx.Exec(fmt.Sprintf("SET LOCAL hnsw.ef_search = %d", efSearch)).Error; err != nil {
|
||||
// Treat as non-fatal: pgvector should always expose this GUC, but if
|
||||
// for any reason it does not we still want the query to run (just
|
||||
// with default recall). We must rollback first because a failed
|
||||
// statement aborts the transaction in PostgreSQL.
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] Failed to set hnsw.ef_search=%d: %v", efSearch, err)
|
||||
return err
|
||||
}
|
||||
// pgvector >= 0.8 supports iterative scan, which keeps pulling more
|
||||
// candidates from HNSW until the post-filter (knowledge_base_id /
|
||||
// knowledge_id / tag_id / is_enabled) yields enough rows. Without it,
|
||||
// HNSW returns at most ef_search candidates and the outer filter may
|
||||
// silently lose recall when the filter is selective.
|
||||
// Best-effort: ignore failure on older pgvector versions.
|
||||
if err := tx.Exec("SET LOCAL hnsw.iterative_scan = strict_order").Error; err != nil {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] hnsw.iterative_scan not available: %v", err)
|
||||
// abort transaction and let the fallback path below handle it.
|
||||
return err
|
||||
}
|
||||
return tx.Raw(querySQL, allVars...).Scan(&embeddingDBList).Error
|
||||
})
|
||||
|
||||
// Fallback: if the transaction failed because of an unsupported GUC (e.g.
|
||||
// older pgvector that doesn't have hnsw.ef_search or hnsw.iterative_scan),
|
||||
// retry the query without the SETs so we still return results.
|
||||
if err != nil && len(embeddingDBList) == 0 &&
|
||||
(strings.Contains(err.Error(), "hnsw.ef_search") ||
|
||||
strings.Contains(err.Error(), "hnsw.iterative_scan")) {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] Retrying vector query without HNSW GUC overrides: %v", err)
|
||||
err = g.db.WithContext(ctx).Raw(querySQL, allVars...).Scan(&embeddingDBList).Error
|
||||
}
|
||||
|
||||
if err == gorm.ErrRecordNotFound {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] No vector matches found that meet threshold %.4f", params.Threshold)
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Vector retrieval failed: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Apply final TopK limit (in case we got more results than needed)
|
||||
if len(embeddingDBList) > int(params.TopK) {
|
||||
embeddingDBList = embeddingDBList[:params.TopK]
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Vector retrieval found %d results", len(embeddingDBList))
|
||||
results := make([]*types.IndexWithScore, len(embeddingDBList))
|
||||
const maxVectorResultLog = 8
|
||||
for i := range embeddingDBList {
|
||||
results[i] = fromDBVectorEmbeddingWithScore(&embeddingDBList[i], types.MatchTypeEmbedding)
|
||||
if i < maxVectorResultLog {
|
||||
logger.GetLogger(ctx).Debugf("[Postgres] Vector search result %d: chunk_id %s, score %.4f",
|
||||
i, results[i].ChunkID, results[i].Score)
|
||||
}
|
||||
}
|
||||
if len(results) > maxVectorResultLog {
|
||||
logger.GetLogger(ctx).Debugf(
|
||||
"[Postgres] Vector search result summary: total=%d logged=%d truncated=%d",
|
||||
len(results), maxVectorResultLog, len(results)-maxVectorResultLog,
|
||||
)
|
||||
}
|
||||
return []*types.RetrieveResult{
|
||||
{
|
||||
Results: results,
|
||||
RetrieverEngineType: types.PostgresRetrieverEngineType,
|
||||
RetrieverType: types.VectorRetrieverType,
|
||||
Error: nil,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// CopyIndices copies index data
|
||||
func (g *pgRepository) CopyIndices(ctx context.Context,
|
||||
sourceKnowledgeBaseID string,
|
||||
sourceToTargetKBIDMap map[string]string,
|
||||
sourceToTargetChunkIDMap map[string]string,
|
||||
targetKnowledgeBaseID string,
|
||||
dimension int,
|
||||
knowledgeType string,
|
||||
) error {
|
||||
logger.GetLogger(ctx).Infof(
|
||||
"[Postgres] Copying indices, source knowledge base: %s, target knowledge base: %s, mapping count: %d",
|
||||
sourceKnowledgeBaseID, targetKnowledgeBaseID, len(sourceToTargetChunkIDMap),
|
||||
)
|
||||
|
||||
if len(sourceToTargetChunkIDMap) == 0 {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] Mapping is empty, no need to copy")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Batch processing parameters
|
||||
batchSize := 500 // Number of records to process per batch
|
||||
offset := 0 // Offset for pagination
|
||||
totalCopied := 0 // Total number of copied records
|
||||
|
||||
for {
|
||||
// Paginated query for source data
|
||||
var sourceVectors []*pgVector
|
||||
if err := g.db.WithContext(ctx).
|
||||
Where("knowledge_base_id = ?", sourceKnowledgeBaseID).
|
||||
Limit(batchSize).
|
||||
Offset(offset).
|
||||
Find(&sourceVectors).Error; err != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to query source index data: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// If no more data, exit the loop
|
||||
if len(sourceVectors) == 0 {
|
||||
if offset == 0 {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] No source index data found")
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
batchCount := len(sourceVectors)
|
||||
logger.GetLogger(ctx).Infof(
|
||||
"[Postgres] Found %d source index data, batch start position: %d",
|
||||
batchCount, offset,
|
||||
)
|
||||
|
||||
// Create target vector index
|
||||
targetVectors := make([]*pgVector, 0, batchCount)
|
||||
for _, sourceVector := range sourceVectors {
|
||||
// Get the mapped target chunk ID
|
||||
targetChunkID, ok := sourceToTargetChunkIDMap[sourceVector.ChunkID]
|
||||
if !ok {
|
||||
logger.GetLogger(ctx).Warnf(
|
||||
"[Postgres] Source chunk %s not found in target chunk mapping, skipping",
|
||||
sourceVector.ChunkID,
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// Get the mapped target knowledge ID
|
||||
targetKnowledgeID, ok := sourceToTargetKBIDMap[sourceVector.KnowledgeID]
|
||||
if !ok {
|
||||
logger.GetLogger(ctx).Warnf(
|
||||
"[Postgres] Source knowledge %s not found in target knowledge mapping, skipping",
|
||||
sourceVector.KnowledgeID,
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle SourceID transformation for generated questions
|
||||
// Generated questions have SourceID format: {chunkID}-{questionID}
|
||||
// Regular chunks have SourceID == ChunkID
|
||||
var targetSourceID string
|
||||
if sourceVector.SourceID == sourceVector.ChunkID {
|
||||
// Regular chunk, use targetChunkID as SourceID
|
||||
targetSourceID = targetChunkID
|
||||
} else if strings.HasPrefix(sourceVector.SourceID, sourceVector.ChunkID+"-") {
|
||||
// This is a generated question, preserve the questionID part
|
||||
questionID := strings.TrimPrefix(sourceVector.SourceID, sourceVector.ChunkID+"-")
|
||||
targetSourceID = fmt.Sprintf("%s-%s", targetChunkID, questionID)
|
||||
} else {
|
||||
// For other complex scenarios, generate new unique SourceID
|
||||
targetSourceID = uuid.New().String()
|
||||
}
|
||||
|
||||
// Create new vector index, copy the content and vector of the source index
|
||||
targetVector := &pgVector{
|
||||
Content: sourceVector.Content,
|
||||
SourceID: targetSourceID, // Handle SourceID transformation properly
|
||||
SourceType: sourceVector.SourceType,
|
||||
ChunkID: targetChunkID, // Update to target chunk ID
|
||||
KnowledgeID: targetKnowledgeID, // Update to target knowledge ID
|
||||
KnowledgeBaseID: targetKnowledgeBaseID, // Update to target knowledge base ID
|
||||
Dimension: sourceVector.Dimension,
|
||||
Embedding: sourceVector.Embedding, // Copy the vector embedding directly, avoid recalculation
|
||||
}
|
||||
|
||||
targetVectors = append(targetVectors, targetVector)
|
||||
}
|
||||
|
||||
// Batch insert target vector index
|
||||
if len(targetVectors) > 0 {
|
||||
if err := g.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{DoNothing: true}).Create(targetVectors).Error; err != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to batch create target index: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
totalCopied += len(targetVectors)
|
||||
logger.GetLogger(ctx).Infof(
|
||||
"[Postgres] Successfully copied batch data, batch size: %d, total copied: %d",
|
||||
len(targetVectors),
|
||||
totalCopied,
|
||||
)
|
||||
}
|
||||
|
||||
// Move to the next batch
|
||||
offset += batchCount
|
||||
|
||||
// If the number of returned records is less than the requested size, it means the last page has been reached
|
||||
if batchCount < batchSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Index copying completed, total copied: %d", totalCopied)
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchUpdateChunkEnabledStatus updates the enabled status of chunks in batch
|
||||
func (g *pgRepository) BatchUpdateChunkEnabledStatus(ctx context.Context, chunkStatusMap map[string]bool) error {
|
||||
if len(chunkStatusMap) == 0 {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] Chunk status map is empty, skipping update")
|
||||
return nil
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Batch updating chunk enabled status, count: %d", len(chunkStatusMap))
|
||||
|
||||
// Group chunks by enabled status for batch updates
|
||||
enabledChunkIDs := make([]string, 0)
|
||||
disabledChunkIDs := make([]string, 0)
|
||||
|
||||
for chunkID, enabled := range chunkStatusMap {
|
||||
if enabled {
|
||||
enabledChunkIDs = append(enabledChunkIDs, chunkID)
|
||||
} else {
|
||||
disabledChunkIDs = append(disabledChunkIDs, chunkID)
|
||||
}
|
||||
}
|
||||
|
||||
// Batch update enabled chunks
|
||||
if len(enabledChunkIDs) > 0 {
|
||||
result := g.db.WithContext(ctx).Model(&pgVector{}).
|
||||
Where("chunk_id IN ?", enabledChunkIDs).
|
||||
Update("is_enabled", true)
|
||||
if result.Error != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to update enabled chunks: %v", result.Error)
|
||||
return result.Error
|
||||
}
|
||||
logger.GetLogger(ctx).
|
||||
Infof("[Postgres] Updated %d chunks to enabled, rows affected: %d", len(enabledChunkIDs), result.RowsAffected)
|
||||
}
|
||||
|
||||
// Batch update disabled chunks
|
||||
if len(disabledChunkIDs) > 0 {
|
||||
result := g.db.WithContext(ctx).Model(&pgVector{}).
|
||||
Where("chunk_id IN ?", disabledChunkIDs).
|
||||
Update("is_enabled", false)
|
||||
if result.Error != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to update disabled chunks: %v", result.Error)
|
||||
return result.Error
|
||||
}
|
||||
logger.GetLogger(ctx).
|
||||
Infof("[Postgres] Updated %d chunks to disabled, rows affected: %d", len(disabledChunkIDs), result.RowsAffected)
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully batch updated chunk enabled status")
|
||||
return nil
|
||||
}
|
||||
|
||||
// BatchUpdateChunkTagID updates the tag ID of chunks in batch
|
||||
func (g *pgRepository) BatchUpdateChunkTagID(ctx context.Context, chunkTagMap map[string]string) error {
|
||||
if len(chunkTagMap) == 0 {
|
||||
logger.GetLogger(ctx).Warnf("[Postgres] Chunk tag map is empty, skipping update")
|
||||
return nil
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Batch updating chunk tag ID, count: %d", len(chunkTagMap))
|
||||
|
||||
// Group chunks by tag ID for batch updates
|
||||
tagGroups := make(map[string][]string)
|
||||
for chunkID, tagID := range chunkTagMap {
|
||||
tagGroups[tagID] = append(tagGroups[tagID], chunkID)
|
||||
}
|
||||
|
||||
// Batch update chunks for each tag ID
|
||||
for tagID, chunkIDs := range tagGroups {
|
||||
result := g.db.WithContext(ctx).Model(&pgVector{}).
|
||||
Where("chunk_id IN ?", chunkIDs).
|
||||
Update("tag_id", tagID)
|
||||
if result.Error != nil {
|
||||
logger.GetLogger(ctx).Errorf("[Postgres] Failed to update chunks with tag_id %s: %v", tagID, result.Error)
|
||||
return result.Error
|
||||
}
|
||||
logger.GetLogger(ctx).
|
||||
Infof("[Postgres] Updated %d chunks to tag_id=%s, rows affected: %d", len(chunkIDs), tagID, result.RowsAffected)
|
||||
}
|
||||
|
||||
logger.GetLogger(ctx).Infof("[Postgres] Successfully batch updated chunk tag ID")
|
||||
return nil
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user