- Added GetCollection endpoint to retrieve collections by ID, including error handling for missing IDs and uninitialized knowledge base. - Updated CreateCollectionConfig to use more descriptive field names for embedding provider and option IDs. - Refactored UpdateSegments function to improve validation and error handling, ensuring document IDs and segment texts are properly checked. - Introduced new request structures for updating votes, scores, and weights for segments, enhancing the API's capabilities. - Removed the outdated vote handling file to streamline the codebase.
517 lines
12 KiB
Go
517 lines
12 KiB
Go
package openapi_test
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/yaoapp/gou/graphrag/types"
|
|
kbtypes "github.com/yaoapp/yao/kb/types"
|
|
"github.com/yaoapp/yao/openapi/kb"
|
|
"github.com/yaoapp/yao/openapi/tests/testutils"
|
|
)
|
|
|
|
func TestProviderConfig_Validation(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
config *kb.ProviderConfig
|
|
expectError bool
|
|
}{
|
|
{
|
|
name: "valid provider config with option_id",
|
|
config: &kb.ProviderConfig{
|
|
ProviderID: "test_provider",
|
|
OptionID: "test_option",
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "valid provider config with direct option",
|
|
config: &kb.ProviderConfig{
|
|
ProviderID: "test_provider",
|
|
Option: &kbtypes.ProviderOption{
|
|
Label: "Test Option",
|
|
Value: "test",
|
|
Description: "Test description",
|
|
Properties: map[string]interface{}{"key": "value"},
|
|
},
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "invalid provider config - empty provider_id",
|
|
config: &kb.ProviderConfig{
|
|
ProviderID: "",
|
|
OptionID: "test_option",
|
|
},
|
|
expectError: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
if tt.config.ProviderID == "" && !tt.expectError {
|
|
t.Errorf("Expected validation to fail for empty provider_id")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBaseUpsertRequest_Validate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
request *kb.BaseUpsertRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "valid base request",
|
|
request: &kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
OptionID: "default",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
OptionID: "default",
|
|
},
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "missing collection_id",
|
|
request: &kb.BaseUpsertRequest{
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
expectError: true,
|
|
errorMsg: "collection_id is required",
|
|
},
|
|
{
|
|
name: "missing chunking provider",
|
|
request: &kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
expectError: true,
|
|
errorMsg: "chunking provider is required",
|
|
},
|
|
{
|
|
name: "missing embedding provider",
|
|
request: &kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
},
|
|
expectError: true,
|
|
errorMsg: "embedding provider is required",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
err := tt.request.Validate()
|
|
|
|
if tt.expectError {
|
|
if err == nil {
|
|
t.Errorf("Expected error but got none")
|
|
} else if err.Error() != tt.errorMsg {
|
|
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
|
|
}
|
|
} else {
|
|
if err != nil {
|
|
t.Errorf("Expected no error but got: %v", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAddFileRequest_Validate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
request *kb.AddFileRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "valid add file request",
|
|
request: &kb.AddFileRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
FileID: "test_file_123",
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "missing file_id",
|
|
request: &kb.AddFileRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
FileID: "",
|
|
},
|
|
expectError: true,
|
|
errorMsg: "file_id is required",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
err := tt.request.Validate()
|
|
|
|
if tt.expectError {
|
|
if err == nil {
|
|
t.Errorf("Expected error but got none")
|
|
} else if err.Error() != tt.errorMsg {
|
|
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
|
|
}
|
|
} else {
|
|
if err != nil {
|
|
t.Errorf("Expected no error but got: %v", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAddTextRequest_Validate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
request *kb.AddTextRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "valid add text request",
|
|
request: &kb.AddTextRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
Text: "This is test text content",
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "missing text",
|
|
request: &kb.AddTextRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
Text: "",
|
|
},
|
|
expectError: true,
|
|
errorMsg: "text is required",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
err := tt.request.Validate()
|
|
|
|
if tt.expectError {
|
|
if err == nil {
|
|
t.Errorf("Expected error but got none")
|
|
} else if err.Error() != tt.errorMsg {
|
|
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
|
|
}
|
|
} else {
|
|
if err != nil {
|
|
t.Errorf("Expected no error but got: %v", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAddURLRequest_Validate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
request *kb.AddURLRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "valid add URL request",
|
|
request: &kb.AddURLRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
URL: "https://example.com/document",
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "missing URL",
|
|
request: &kb.AddURLRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
URL: "",
|
|
},
|
|
expectError: true,
|
|
errorMsg: "url is required",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
err := tt.request.Validate()
|
|
|
|
if tt.expectError {
|
|
if err == nil {
|
|
t.Errorf("Expected error but got none")
|
|
} else if err.Error() != tt.errorMsg {
|
|
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
|
|
}
|
|
} else {
|
|
if err != nil {
|
|
t.Errorf("Expected no error but got: %v", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAddSegmentsRequest_Validate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
request *kb.AddSegmentsRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "valid add segments request",
|
|
request: &kb.AddSegmentsRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
DocID: "test_doc_123",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
SegmentTexts: []types.SegmentText{
|
|
{Text: "First segment"},
|
|
{Text: "Second segment"},
|
|
},
|
|
},
|
|
expectError: false,
|
|
},
|
|
{
|
|
name: "missing segment_texts",
|
|
request: &kb.AddSegmentsRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
DocID: "test_doc_123",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
SegmentTexts: []types.SegmentText{},
|
|
},
|
|
expectError: true,
|
|
errorMsg: "segment_texts is required",
|
|
},
|
|
{
|
|
name: "missing doc_id",
|
|
request: &kb.AddSegmentsRequest{
|
|
BaseUpsertRequest: kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
},
|
|
},
|
|
SegmentTexts: []types.SegmentText{
|
|
{Text: "Test segment"},
|
|
},
|
|
},
|
|
expectError: true,
|
|
errorMsg: "doc_id is required for AddSegments operation",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
err := tt.request.Validate()
|
|
|
|
if tt.expectError {
|
|
if err == nil {
|
|
t.Errorf("Expected error but got none")
|
|
} else if err.Error() != tt.errorMsg {
|
|
t.Errorf("Expected error message '%s', got '%s'", tt.errorMsg, err.Error())
|
|
}
|
|
} else {
|
|
if err != nil {
|
|
t.Errorf("Expected no error but got: %v", err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestUpdateSegmentsRequest_Structure(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
request *kb.UpdateSegmentsRequest
|
|
expectValid bool
|
|
}{
|
|
{
|
|
name: "valid update segments request",
|
|
request: &kb.UpdateSegmentsRequest{
|
|
SegmentTexts: []types.SegmentText{
|
|
{ID: "segment_1", Text: "Updated segment"},
|
|
},
|
|
},
|
|
expectValid: true,
|
|
},
|
|
{
|
|
name: "empty segment_texts",
|
|
request: &kb.UpdateSegmentsRequest{
|
|
SegmentTexts: []types.SegmentText{},
|
|
},
|
|
expectValid: false,
|
|
},
|
|
{
|
|
name: "multiple segments",
|
|
request: &kb.UpdateSegmentsRequest{
|
|
SegmentTexts: []types.SegmentText{
|
|
{ID: "segment_1", Text: "First updated segment"},
|
|
{ID: "segment_2", Text: "Second updated segment"},
|
|
},
|
|
},
|
|
expectValid: true,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
// Test basic structure validation
|
|
if tt.expectValid {
|
|
if len(tt.request.SegmentTexts) == 0 {
|
|
t.Errorf("Expected valid request to have segment_texts")
|
|
}
|
|
} else {
|
|
if len(tt.request.SegmentTexts) > 0 {
|
|
t.Errorf("Expected invalid request to have empty segment_texts")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestToUpsertOptions_WithEnvironment(t *testing.T) {
|
|
// Initialize test environment
|
|
serverURL := testutils.Prepare(t)
|
|
defer testutils.Clean()
|
|
|
|
t.Logf("Test server running at: %s", serverURL)
|
|
|
|
// Test with proper environment setup and direct options to avoid provider lookup
|
|
request := &kb.BaseUpsertRequest{
|
|
CollectionID: "test_collection",
|
|
DocID: "test_doc",
|
|
Metadata: map[string]interface{}{"source": "test"},
|
|
Chunking: &kb.ProviderConfig{
|
|
ProviderID: "chunking_provider",
|
|
Option: &kbtypes.ProviderOption{
|
|
Label: "Test Chunking",
|
|
Value: "test",
|
|
Description: "Test chunking option",
|
|
Properties: map[string]interface{}{"chunk_size": 1000},
|
|
},
|
|
},
|
|
Embedding: &kb.ProviderConfig{
|
|
ProviderID: "embedding_provider",
|
|
Option: &kbtypes.ProviderOption{
|
|
Label: "Test Embedding",
|
|
Value: "test",
|
|
Description: "Test embedding option",
|
|
Properties: map[string]interface{}{"model": "test-model"},
|
|
},
|
|
},
|
|
}
|
|
|
|
t.Run("no parameters", func(t *testing.T) {
|
|
options, err := request.ToUpsertOptions()
|
|
if err != nil {
|
|
t.Logf("ToUpsertOptions failed (expected with test environment): %v", err)
|
|
// This is expected to fail in test environment due to missing actual providers
|
|
// But we can verify the basic structure is being built correctly
|
|
return
|
|
}
|
|
|
|
// If it doesn't fail, verify the basic structure
|
|
if options.CollectionID != "test_collection" {
|
|
t.Errorf("Expected CollectionID to be 'test_collection', got '%s'", options.CollectionID)
|
|
}
|
|
if options.DocID != "test_doc" {
|
|
t.Errorf("Expected DocID to be 'test_doc', got '%s'", options.DocID)
|
|
}
|
|
})
|
|
|
|
t.Run("with filename and contentType", func(t *testing.T) {
|
|
options, err := request.ToUpsertOptions("test.pdf", "application/pdf")
|
|
if err != nil {
|
|
t.Logf("ToUpsertOptions with file info failed (expected with test environment): %v", err)
|
|
// This is expected to fail in test environment due to missing actual providers
|
|
return
|
|
}
|
|
|
|
// If it doesn't fail, verify the basic structure
|
|
if options.CollectionID != "test_collection" {
|
|
t.Errorf("Expected CollectionID to be 'test_collection', got '%s'", options.CollectionID)
|
|
}
|
|
})
|
|
}
|