Rogeeandrogee
d948222ac6
* fix(captain): restore inbox takeover and KB citations * fix(captain): harden grounded citations and smoke seed --------- Co-authored-by: Rogee <rogee@ipao.vip>
343 lines
14 KiB
Go
343 lines
14 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/pgvector/pgvector-go"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/gorm"
|
|
|
|
"github.com/gochat/gochat/internal/model"
|
|
)
|
|
|
|
func createCaptainResponseTestAssistant(t *testing.T, db *gorm.DB) *model.CaptainAssistant {
|
|
t.Helper()
|
|
account := &model.Account{Name: "Captain response test", Active: true}
|
|
require.NoError(t, db.Create(account).Error)
|
|
assistant := &model.CaptainAssistant{AccountID: account.ID, Name: "Test assistant"}
|
|
require.NoError(t, db.Create(assistant).Error)
|
|
return assistant
|
|
}
|
|
|
|
// zeroEmbedding returns a 1536-dimensional zero vector suitable for SQLite compatibility.
|
|
// SQLite cannot scan an empty string as pgvector; providing an explicit zero vector
|
|
// ensures the embedding column is always populated with a parseable value.
|
|
func zeroEmbedding() pgvector.Vector {
|
|
dims := make([]float32, 1536)
|
|
return pgvector.NewVector(dims)
|
|
}
|
|
|
|
// createTestResponse builds a minimal valid CaptainAssistantResponse with a zero embedding.
|
|
func createTestResponse(accountID, assistantID uint, question, answer string) *model.CaptainAssistantResponse {
|
|
return &model.CaptainAssistantResponse{
|
|
AccountID: accountID,
|
|
AssistantID: assistantID,
|
|
Question: question,
|
|
Answer: answer,
|
|
Status: model.ResponseStatusApproved,
|
|
Edited: false,
|
|
Embedding: zeroEmbedding(),
|
|
}
|
|
}
|
|
|
|
// createTestResponseWithDocument builds a response linked to a documentable with a zero embedding.
|
|
func createTestResponseWithDocument(accountID, assistantID uint, documentableID uint, documentableType, question, answer string) *model.CaptainAssistantResponse {
|
|
resp := createTestResponse(accountID, assistantID, question, answer)
|
|
resp.DocumentableID = &documentableID
|
|
resp.DocumentableType = documentableType
|
|
return resp
|
|
}
|
|
|
|
// ========== 1. Create ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_Create(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
resp := createTestResponse(1, 10, "What is GoChat?", "GoChat is an open-source chat platform.")
|
|
err := repo.Create(context.Background(), resp)
|
|
require.NoError(t, err)
|
|
assert.NotZero(t, resp.ID, "ID should be set after Create")
|
|
assert.Equal(t, uint(1), resp.AccountID)
|
|
assert.Equal(t, uint(10), resp.AssistantID)
|
|
assert.Equal(t, "What is GoChat?", resp.Question)
|
|
assert.Equal(t, "GoChat is an open-source chat platform.", resp.Answer)
|
|
assert.Equal(t, model.ResponseStatusApproved, resp.Status)
|
|
}
|
|
|
|
// ========== 2. GetByID ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_GetByID(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
resp := createTestResponse(1, 10, "GetByID question", "GetByID answer")
|
|
err := repo.Create(context.Background(), resp)
|
|
require.NoError(t, err)
|
|
|
|
found, err := repo.GetByID(context.Background(), resp.ID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, resp.ID, found.ID)
|
|
assert.Equal(t, "GetByID question", found.Question)
|
|
assert.Equal(t, uint(10), found.AssistantID)
|
|
}
|
|
|
|
// ========== 3. GetByID Not Found ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_GetByID_NotFound(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
found, err := repo.GetByID(context.Background(), 9999)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, found)
|
|
}
|
|
|
|
// ========== 4. Update ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_Update(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
resp := createTestResponse(1, 10, "Original question", "Original answer")
|
|
err := repo.Create(context.Background(), resp)
|
|
require.NoError(t, err)
|
|
|
|
// Modify fields and update
|
|
resp.Question = "Updated question"
|
|
resp.Answer = "Updated answer"
|
|
resp.Status = model.ResponseStatusPending
|
|
resp.Edited = true
|
|
|
|
err = repo.Update(context.Background(), resp)
|
|
require.NoError(t, err)
|
|
|
|
found, err := repo.GetByID(context.Background(), resp.ID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "Updated question", found.Question)
|
|
assert.Equal(t, "Updated answer", found.Answer)
|
|
assert.Equal(t, model.ResponseStatusPending, found.Status)
|
|
assert.True(t, found.Edited)
|
|
}
|
|
|
|
// ========== 5. Delete ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_Delete(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
resp := createTestResponse(1, 10, "Delete question", "Delete answer")
|
|
err := repo.Create(context.Background(), resp)
|
|
require.NoError(t, err)
|
|
|
|
err = repo.Delete(context.Background(), resp.ID)
|
|
require.NoError(t, err)
|
|
|
|
// After deletion, GetByID should return error
|
|
found, err := repo.GetByID(context.Background(), resp.ID)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, found)
|
|
}
|
|
|
|
// ========== 6. ListByAssistant ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_ListByAssistant(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
// Create 3 responses under assistant 10, account 1
|
|
for i := 0; i < 3; i++ {
|
|
r := createTestResponse(1, 10, "Q10-"+string(rune('A'+i)), "A10-"+string(rune('A'+i)))
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
}
|
|
// Create 2 responses under assistant 20, account 1 (should not appear)
|
|
for i := 0; i < 2; i++ {
|
|
r := createTestResponse(1, 20, "Q20-"+string(rune('A'+i)), "A20-"+string(rune('A'+i)))
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
}
|
|
|
|
responses, count, err := repo.ListByAssistant(context.Background(), 10, 0, 10)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(3), count)
|
|
assert.Len(t, responses, 3)
|
|
|
|
for _, r := range responses {
|
|
assert.Equal(t, uint(10), r.AssistantID)
|
|
}
|
|
}
|
|
|
|
// ========== 7. ListByAssistant with Pagination ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_ListByAssistant_Pagination(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
// Create 5 responses under assistant 10
|
|
for i := 0; i < 5; i++ {
|
|
r := createTestResponse(1, 10, "PagQ-"+string(rune('A'+i)), "PagA-"+string(rune('A'+i)))
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
}
|
|
|
|
// First page: offset=0, limit=2
|
|
responses, count, err := repo.ListByAssistant(context.Background(), 10, 0, 2)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(5), count, "total count should be 5 regardless of pagination")
|
|
assert.Len(t, responses, 2, "first page should return 2 items")
|
|
|
|
// Second page: offset=2, limit=2
|
|
responses2, count2, err := repo.ListByAssistant(context.Background(), 10, 2, 2)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(5), count2)
|
|
assert.Len(t, responses2, 2, "second page should return 2 items")
|
|
|
|
// Third page: offset=4, limit=2
|
|
responses3, count3, err := repo.ListByAssistant(context.Background(), 10, 4, 2)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(5), count3)
|
|
assert.Len(t, responses3, 1, "last page should return 1 item")
|
|
}
|
|
|
|
// ========== 8. ListByDocument ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_ListByDocument(t *testing.T) {
|
|
skipIfSQLite(t)
|
|
db := setupTestDB(t, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
|
|
// Create 2 responses linked to document (ID=100, type="CaptainDocument")
|
|
for i := 0; i < 2; i++ {
|
|
r := createTestResponseWithDocument(1, 10, 100, "CaptainDocument",
|
|
"DocQ-"+string(rune('A'+i)), "DocA-"+string(rune('A'+i)))
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
}
|
|
// Create 1 response linked to assistant (documentableType="CaptainAssistant") — should not appear
|
|
r := createTestResponseWithDocument(1, 10, 50, "CaptainAssistant",
|
|
"AssistQ", "AssistA")
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
|
|
responses, count, err := repo.ListByDocument(context.Background(), 100, "CaptainDocument", 0, 10)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int64(2), count)
|
|
assert.Len(t, responses, 2)
|
|
|
|
for _, resp := range responses {
|
|
assert.Equal(t, uint(100), *resp.DocumentableID)
|
|
assert.Equal(t, "CaptainDocument", resp.DocumentableType)
|
|
}
|
|
}
|
|
|
|
// ========== 9. SimilaritySearch (PG only) ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_SimilaritySearch(t *testing.T) {
|
|
skipIfSQLite(t) // pgvector requires PostgreSQL
|
|
|
|
db := setupTestDB(t, &model.CaptainAssistant{}, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
assistant := createCaptainResponseTestAssistant(t, db)
|
|
|
|
// Create approved responses under the test assistant.
|
|
for i := 0; i < 3; i++ {
|
|
r := createTestResponse(assistant.AccountID, assistant.ID, "SimilarQ-"+string(rune('A'+i)), "SimilarA-"+string(rune('A'+i)))
|
|
r.Status = model.ResponseStatusApproved
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
}
|
|
// Create one rejected response — should be excluded from similarity search
|
|
r := createTestResponse(assistant.AccountID, assistant.ID, "RejectedQ", "RejectedA")
|
|
r.Status = model.ResponseStatusRejected
|
|
require.NoError(t, repo.Create(context.Background(), r))
|
|
|
|
// Build a dummy 1536-dimensional embedding vector for search
|
|
dims := make([]float32, 1536)
|
|
for i := range dims {
|
|
dims[i] = 0.01
|
|
}
|
|
embedding := pgvector.NewVector(dims)
|
|
|
|
results, err := repo.SimilaritySearch(context.Background(), assistant.ID, embedding, 5)
|
|
require.NoError(t, err)
|
|
// Only approved responses (3) should be returned
|
|
assert.Len(t, results, 3)
|
|
for _, resp := range results {
|
|
assert.Equal(t, model.ResponseStatusApproved, resp.Status)
|
|
assert.Equal(t, assistant.ID, resp.AssistantID)
|
|
}
|
|
}
|
|
|
|
// ========== 10. SearchByEmbedding (alias for SimilaritySearch, PG only) ==========
|
|
|
|
func TestCaptainAssistantResponseRepo_SearchByEmbedding(t *testing.T) {
|
|
skipIfSQLite(t) // pgvector requires PostgreSQL
|
|
|
|
db := setupTestDB(t, &model.CaptainAssistant{}, &model.CaptainDocument{}, &model.CaptainAssistantResponse{})
|
|
repo := NewCaptainAssistantResponseRepo(db)
|
|
assistant := createCaptainResponseTestAssistant(t, db)
|
|
otherAssistant := &model.CaptainAssistant{AccountID: assistant.AccountID, Name: "Other assistant"}
|
|
require.NoError(t, db.Create(otherAssistant).Error)
|
|
otherAccount := &model.Account{Name: "Other account", Active: true}
|
|
require.NoError(t, db.Create(otherAccount).Error)
|
|
foreignAssistant := &model.CaptainAssistant{AccountID: otherAccount.ID, Name: "Foreign assistant"}
|
|
require.NoError(t, db.Create(foreignAssistant).Error)
|
|
|
|
document := &model.CaptainDocument{AccountID: assistant.AccountID, AssistantID: assistant.ID, Name: "Source", ExternalLink: "https://example.com/knowledge"}
|
|
crossAssistantDocument := &model.CaptainDocument{AccountID: assistant.AccountID, AssistantID: otherAssistant.ID, Name: "Wrong assistant", ExternalLink: "https://example.com/wrong-assistant"}
|
|
crossAccountDocument := &model.CaptainDocument{AccountID: otherAccount.ID, AssistantID: foreignAssistant.ID, Name: "Wrong account", ExternalLink: "https://example.com/wrong-account"}
|
|
require.NoError(t, db.Create(document).Error)
|
|
require.NoError(t, db.Create(crossAssistantDocument).Error)
|
|
require.NoError(t, db.Create(crossAccountDocument).Error)
|
|
|
|
responses := []*model.CaptainAssistantResponse{
|
|
createTestResponseWithDocument(assistant.AccountID, assistant.ID, document.ID, "Captain::Document", "Valid source", "Valid answer"),
|
|
createTestResponse(assistant.AccountID, assistant.ID, "No source", "Compatible answer"),
|
|
createTestResponseWithDocument(assistant.AccountID, assistant.ID, crossAssistantDocument.ID, "Captain::Document", "Wrong assistant source", "Must not cite"),
|
|
createTestResponseWithDocument(assistant.AccountID, assistant.ID, crossAccountDocument.ID, "Captain::Document", "Wrong account source", "Must not cite"),
|
|
}
|
|
for _, response := range responses {
|
|
require.NoError(t, repo.Create(context.Background(), response))
|
|
}
|
|
|
|
nullEmbedding := createTestResponse(assistant.AccountID, assistant.ID, "Missing embedding", "Must not retrieve")
|
|
require.NoError(t, repo.Create(context.Background(), nullEmbedding))
|
|
require.NoError(t, db.Model(nullEmbedding).UpdateColumn("embedding", gorm.Expr("NULL")).Error)
|
|
corruptAccount := createTestResponse(otherAccount.ID, assistant.ID, "Wrong response account", "Must not retrieve")
|
|
require.NoError(t, repo.Create(context.Background(), corruptAccount))
|
|
allNull := createTestResponse(otherAssistant.AccountID, otherAssistant.ID, "Only null", "Must not retrieve")
|
|
require.NoError(t, repo.Create(context.Background(), allNull))
|
|
require.NoError(t, db.Model(allNull).UpdateColumn("embedding", gorm.Expr("NULL")).Error)
|
|
|
|
dims := make([]float32, 1536)
|
|
for i := range dims {
|
|
dims[i] = 0.02
|
|
}
|
|
embedding := pgvector.NewVector(dims)
|
|
|
|
results, err := repo.SearchByEmbedding(context.Background(), assistant.ID, embedding, 10)
|
|
require.NoError(t, err)
|
|
assert.Len(t, results, 4)
|
|
for _, resp := range results {
|
|
assert.Equal(t, assistant.ID, resp.AssistantID)
|
|
assert.Equal(t, assistant.AccountID, resp.AccountID)
|
|
assert.NotEqual(t, nullEmbedding.ID, resp.ID)
|
|
assert.NotEqual(t, corruptAccount.ID, resp.ID)
|
|
}
|
|
sources := make([]string, 0, len(results))
|
|
for _, result := range results {
|
|
sources = append(sources, result.SourceURL)
|
|
}
|
|
assert.Contains(t, sources, document.ExternalLink)
|
|
assert.NotContains(t, sources, crossAssistantDocument.ExternalLink)
|
|
assert.NotContains(t, sources, crossAccountDocument.ExternalLink)
|
|
|
|
empty, err := repo.SearchByEmbedding(context.Background(), otherAssistant.ID, embedding, 5)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, empty, "an assistant with only NULL embeddings must return no grounded results")
|
|
}
|